Compare commits
5
Commits
55eb87a983
...
75b9334a0b
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
75b9334a0b | ||
|
|
c894a11589 | ||
|
|
3ddd8f25ac | ||
|
|
1820561f2a | ||
|
|
04cff91509 |
@@ -5,6 +5,7 @@ This file provides guidance when working with code in this repository.
|
|||||||
## Development Commands
|
## Development Commands
|
||||||
|
|
||||||
### Build and Run
|
### Build and Run
|
||||||
|
- `cargo build` - Build Galaxy and verify feature and bug-fix changes
|
||||||
- `cargo run` - Build and run Warp locally
|
- `cargo run` - Build and run Warp locally
|
||||||
- `cargo bundle --bin warp` - Bundle the main app
|
- `cargo bundle --bin warp` - Bundle the main app
|
||||||
- `make run` - Build and run `galaxy-oss`
|
- `make run` - Build and run `galaxy-oss`
|
||||||
@@ -28,8 +29,9 @@ Environment variables:
|
|||||||
- `WS_SERVER_URL` - WebSocket endpoint (default: `ws://localhost:8080/graphql/v2`)
|
- `WS_SERVER_URL` - WebSocket endpoint (default: `ws://localhost:8080/graphql/v2`)
|
||||||
|
|
||||||
### Testing
|
### Testing
|
||||||
- Before launching `galaxy-oss`, check for existing app instances and close the prior development instance so only one app instance runs at a time. Distinguish the app process from its `terminal-server` helper; do not count the helper as a duplicate app. Leave the newly launched instance running for the user.
|
- This development machine is CLI-only and has no GUI. The user performs GUI and most behavioral testing on a remote machine.
|
||||||
- During interactive feature and bug-fix verification, make the requested code changes first and run only `cargo run --bin galaxy-oss` for the user to verify. Leave the app running for the user; do not stop it or treat the command timeout as a failure. Do not run `cargo fmt`, `cargo check`, `cargo test`, `cargo clippy`, presubmit, or other validation commands until the user confirms the behavior. Run formatting, tests, and linting only as the final cleanup step after interactive approval.
|
- Verify feature and bug-fix changes with `cargo build` only. Do not use `cargo run`, `make run`, or launch the app to verify results, and do not wait for interactive approval before completing the work.
|
||||||
|
- After `cargo build` succeeds and all requested work and final review are complete, commit the changes and push the current branch to both configured remotes, `origin` and `samsung`, using ordinary non-force pushes. If the build fails, report the blocker and do not push.
|
||||||
- `cargo nextest run --no-fail-fast --workspace --exclude command-signatures-v2` - Run tests with nextest
|
- `cargo nextest run --no-fail-fast --workspace --exclude command-signatures-v2` - Run tests with nextest
|
||||||
- `cargo nextest run -p galaxy_completer --features v2` - Run completer tests with v2 features
|
- `cargo nextest run -p galaxy_completer --features v2` - Run completer tests with v2 features
|
||||||
- `cargo test --doc` - Run doc tests
|
- `cargo test --doc` - Run doc tests
|
||||||
@@ -143,7 +145,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
|
||||||
|
|
||||||
|
|||||||
@@ -14,6 +14,10 @@ use tokio::sync::Mutex;
|
|||||||
use vec1::vec1;
|
use vec1::vec1;
|
||||||
|
|
||||||
use crate::settings::{AISettings, AISettingsChangedEvent};
|
use crate::settings::{AISettings, AISettingsChangedEvent};
|
||||||
|
use crate::ai::provider::authentication::{
|
||||||
|
BedrockAuthConnection, connection_revision, credentials_are_current,
|
||||||
|
};
|
||||||
|
use crate::ai::provider::client::BedrockClientConfig;
|
||||||
use crate::terminal::event::{AfterBlockCompletedEvent, BlockType, UserBlockCompleted};
|
use crate::terminal::event::{AfterBlockCompletedEvent, BlockType, UserBlockCompleted};
|
||||||
use crate::terminal::model_events::{ModelEvent, ModelEventDispatcher};
|
use crate::terminal::model_events::{ModelEvent, ModelEventDispatcher};
|
||||||
use crate::workspaces::user_workspaces::{UserWorkspaces, UserWorkspacesEvent};
|
use crate::workspaces::user_workspaces::{UserWorkspaces, UserWorkspacesEvent};
|
||||||
@@ -56,7 +60,7 @@ fn aws_profile_reference_for_message(profile: &str, capitalize_first_word: bool)
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn user_facing_aws_credentials_error_message(err: &CredentialsError, profile: &str) -> String {
|
pub(crate) fn user_facing_aws_credentials_error_message(err: &CredentialsError, profile: &str) -> String {
|
||||||
match err {
|
match err {
|
||||||
CredentialsError::CredentialsNotLoaded(_) => format!(
|
CredentialsError::CredentialsNotLoaded(_) => format!(
|
||||||
"AWS credentials were not found for {}. Log in with the AWS CLI or update your AWS credentials configuration, then refresh.",
|
"AWS credentials were not found for {}. Log in with the AWS CLI or update your AWS credentials configuration, then refresh.",
|
||||||
@@ -129,15 +133,17 @@ fn aws_credentials_state_for_error(err: LoadAwsCredentialsError) -> AwsCredentia
|
|||||||
/// Loads AWS credentials from the AWS SDK.
|
/// Loads AWS credentials from the AWS SDK.
|
||||||
///
|
///
|
||||||
/// # Arguments
|
/// # Arguments
|
||||||
/// * `profile` - AWS profile name. If empty, uses the default AWS SDK behavior
|
/// * `connection` - Effective AWS connection. If its profile is empty, uses the default AWS SDK behavior
|
||||||
/// (checks AWS_PROFILE env var, then uses "default").
|
/// (checks AWS_PROFILE env var, then uses "default").
|
||||||
pub async fn load_aws_credentials_from_sdk(
|
pub async fn load_aws_credentials_from_sdk(
|
||||||
profile: &str,
|
connection: &BedrockClientConfig,
|
||||||
) -> Result<AwsCredentials, LoadAwsCredentialsError> {
|
) -> Result<AwsCredentials, LoadAwsCredentialsError> {
|
||||||
let region_provider = aws_config::meta::region::RegionProviderChain::default_provider();
|
let profile = &connection.profile;
|
||||||
let loader =
|
let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest());
|
||||||
aws_config::defaults(aws_config::BehaviorVersion::latest()).region(region_provider);
|
if !connection.region.is_empty() {
|
||||||
let loader = if profile.trim().is_empty() || profile.trim() == "default" {
|
loader = loader.region(aws_config::Region::new(connection.region.clone()));
|
||||||
|
}
|
||||||
|
let loader = if profile.trim().is_empty() {
|
||||||
loader // Let AWS SDK use its default behavior
|
loader // Let AWS SDK use its default behavior
|
||||||
} else {
|
} else {
|
||||||
loader.profile_name(profile)
|
loader.profile_name(profile)
|
||||||
@@ -150,14 +156,15 @@ pub async fn load_aws_credentials_from_sdk(
|
|||||||
|
|
||||||
let creds = provider.provide_credentials().await.map_err(|e| {
|
let creds = provider.provide_credentials().await.map_err(|e| {
|
||||||
let message = user_facing_aws_credentials_error_message(&e, profile);
|
let message = user_facing_aws_credentials_error_message(&e, profile);
|
||||||
log::warn!("{e}");
|
log::warn!("AWS credential reload failed: {message}");
|
||||||
// TODO(isaiah): turn this full SDK dump back down to debug once we've resolved
|
|
||||||
// the current customer-facing AWS credential issue and no longer need prod-visible
|
|
||||||
// provider internals for support debugging.
|
|
||||||
log::warn!("{e:#?}");
|
|
||||||
log::info!("AWS credential load failure message shown to user: {message}");
|
log::info!("AWS credential load failure message shown to user: {message}");
|
||||||
LoadAwsCredentialsError::CredentialsLoadFailed(message)
|
LoadAwsCredentialsError::CredentialsLoadFailed(message)
|
||||||
})?;
|
})?;
|
||||||
|
if !credentials_are_current(creds.expiry(), SystemTime::now()) {
|
||||||
|
return Err(LoadAwsCredentialsError::CredentialsLoadFailed(
|
||||||
|
"AWS credentials have expired. Refresh your AWS login.".to_owned(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
Ok(AwsCredentials::new(
|
Ok(AwsCredentials::new(
|
||||||
creds.access_key_id().to_string(),
|
creds.access_key_id().to_string(),
|
||||||
@@ -198,8 +205,10 @@ impl AwsCredentialRefresher for ApiKeyManager {
|
|||||||
..
|
..
|
||||||
}) = event
|
}) = event
|
||||||
{
|
{
|
||||||
let auth_command = &AISettings::as_ref(ctx).bedrock_auth_refresh_command;
|
let auth_command = BedrockAuthConnection::current(ctx).command;
|
||||||
if command.trim().starts_with(auth_command.trim()) {
|
if !auth_command.trim().is_empty()
|
||||||
|
&& command.trim().starts_with(auth_command.trim())
|
||||||
|
{
|
||||||
log::debug!("Detected AWS auth command completion, refreshing credentials");
|
log::debug!("Detected AWS auth command completion, refreshing credentials");
|
||||||
drop(refresh_aws_credentials(manager, ctx));
|
drop(refresh_aws_credentials(manager, ctx));
|
||||||
}
|
}
|
||||||
@@ -225,6 +234,8 @@ impl AwsCredentialRefresher for ApiKeyManager {
|
|||||||
if matches!(
|
if matches!(
|
||||||
event,
|
event,
|
||||||
AISettingsChangedEvent::BedrockProfile { .. }
|
AISettingsChangedEvent::BedrockProfile { .. }
|
||||||
|
| AISettingsChangedEvent::BedrockAuthMethod { .. }
|
||||||
|
| AISettingsChangedEvent::BedrockRegion { .. }
|
||||||
| AISettingsChangedEvent::BedrockAuthRefreshCommand { .. }
|
| AISettingsChangedEvent::BedrockAuthRefreshCommand { .. }
|
||||||
| AISettingsChangedEvent::BedrockEnabled { .. }
|
| AISettingsChangedEvent::BedrockEnabled { .. }
|
||||||
) {
|
) {
|
||||||
@@ -266,15 +277,28 @@ fn refresh_aws_credentials_local_chain(
|
|||||||
return Box::pin(async { Ok(()) });
|
return Box::pin(async { Ok(()) });
|
||||||
}
|
}
|
||||||
|
|
||||||
let profile = (*AISettings::as_ref(ctx).bedrock_profile).clone();
|
let connection = BedrockAuthConnection::current(ctx);
|
||||||
|
if connection.config.auth_method == crate::settings::BedrockAuthMethod::StaticKeys {
|
||||||
|
manager.set_aws_credentials_state(AwsCredentialsState::Missing, ctx);
|
||||||
|
return Box::pin(async { Ok(()) });
|
||||||
|
}
|
||||||
|
let config = connection.config.clone();
|
||||||
|
let revision = connection_revision(ctx);
|
||||||
|
|
||||||
manager.set_aws_credentials_state(AwsCredentialsState::Refreshing, ctx);
|
manager.set_aws_credentials_state(AwsCredentialsState::Refreshing, ctx);
|
||||||
|
|
||||||
let (tx, rx) = channel();
|
let (tx, rx) = channel();
|
||||||
// credential fetch from aws cli's disk cache
|
// credential fetch from aws cli's disk cache
|
||||||
let _ = ctx.spawn(
|
let _ = ctx.spawn(
|
||||||
async move { load_aws_credentials_from_sdk(&profile).await },
|
async move { load_aws_credentials_from_sdk(&config).await },
|
||||||
move |manager, result, ctx| {
|
move |manager, result, ctx| {
|
||||||
|
if connection_revision(ctx) != revision
|
||||||
|
|| BedrockAuthConnection::current(ctx) != connection
|
||||||
|
|| manager.aws_credentials_refresh_strategy() != AwsCredentialsRefreshStrategy::LocalChain
|
||||||
|
{
|
||||||
|
let _ = tx.send(Err("AWS connection changed during credential reload".to_owned()));
|
||||||
|
return;
|
||||||
|
}
|
||||||
let (new_state, tx_result) = match result {
|
let (new_state, tx_result) = match result {
|
||||||
Ok(credentials) => (
|
Ok(credentials) => (
|
||||||
AwsCredentialsState::Loaded {
|
AwsCredentialsState::Loaded {
|
||||||
|
|||||||
@@ -4439,7 +4439,7 @@ impl AIBlock {
|
|||||||
}
|
}
|
||||||
|
|
||||||
let ai_settings = AISettings::as_ref(ctx);
|
let ai_settings = AISettings::as_ref(ctx);
|
||||||
let login_command = ai_settings.bedrock_auth_refresh_command.value().clone();
|
let login_command = crate::ai::provider::authentication::BedrockAuthConnection::current(ctx).command;
|
||||||
// Recreating a historical error view must not restart authentication.
|
// Recreating a historical error view must not restart authentication.
|
||||||
let auto_login_enabled = attempt_auto_login && *ai_settings.bedrock_auto_login.value();
|
let auto_login_enabled = attempt_auto_login && *ai_settings.bedrock_auto_login.value();
|
||||||
|
|
||||||
@@ -4452,7 +4452,7 @@ impl AIBlock {
|
|||||||
|
|
||||||
let model_name = model_name.clone();
|
let model_name = model_name.clone();
|
||||||
let view = ctx.add_typed_action_view(|ctx| {
|
let view = ctx.add_typed_action_view(|ctx| {
|
||||||
AwsBedrockCredentialsErrorView::new(model_name, login_command, auto_login_enabled, ctx)
|
AwsBedrockCredentialsErrorView::new(model_name, login_command, ctx)
|
||||||
});
|
});
|
||||||
|
|
||||||
// Subscribe to events from the view and emit AIBlockEvents directly
|
// Subscribe to events from the view and emit AIBlockEvents directly
|
||||||
|
|||||||
@@ -6610,7 +6610,11 @@ impl BlocklistAIController {
|
|||||||
let prepared = match result {
|
let prepared = match result {
|
||||||
Ok(prepared) => prepared,
|
Ok(prepared) => prepared,
|
||||||
Err(error) => {
|
Err(error) => {
|
||||||
self.fail_provider_startup(conversation_id, stream_id, error.to_string(), ctx);
|
let renderable = crate::ai::provider::authentication::provider_startup_error(
|
||||||
|
&error,
|
||||||
|
slot.response_stream.as_ref(ctx).model_id(),
|
||||||
|
);
|
||||||
|
self.fail_provider_startup_with_error(conversation_id, stream_id, renderable, ctx);
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@@ -8916,6 +8920,21 @@ impl BlocklistAIController {
|
|||||||
stream_id: ResponseStreamId,
|
stream_id: ResponseStreamId,
|
||||||
message: String,
|
message: String,
|
||||||
ctx: &mut ModelContext<Self>,
|
ctx: &mut ModelContext<Self>,
|
||||||
|
) {
|
||||||
|
self.fail_provider_startup_with_error(
|
||||||
|
conversation_id,
|
||||||
|
stream_id,
|
||||||
|
RenderableAIError::other(message, false),
|
||||||
|
ctx,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
fn fail_provider_startup_with_error(
|
||||||
|
&mut self,
|
||||||
|
conversation_id: AIConversationId,
|
||||||
|
stream_id: ResponseStreamId,
|
||||||
|
error: RenderableAIError,
|
||||||
|
ctx: &mut ModelContext<Self>,
|
||||||
) {
|
) {
|
||||||
let response_stream = self
|
let response_stream = self
|
||||||
.active_provider_runs
|
.active_provider_runs
|
||||||
@@ -8925,14 +8944,10 @@ impl BlocklistAIController {
|
|||||||
let Some(response_stream) = response_stream else {
|
let Some(response_stream) = response_stream else {
|
||||||
return;
|
return;
|
||||||
};
|
};
|
||||||
|
let message = error.to_string();
|
||||||
BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| {
|
BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| {
|
||||||
history_model.mark_response_stream_completed_with_error(
|
history_model.mark_response_stream_completed_with_error(
|
||||||
RenderableAIError::Other {
|
error,
|
||||||
error_message: message.clone(),
|
|
||||||
will_attempt_resume: false,
|
|
||||||
waiting_for_network: false,
|
|
||||||
is_user_error: false,
|
|
||||||
},
|
|
||||||
false,
|
false,
|
||||||
&stream_id,
|
&stream_id,
|
||||||
conversation_id,
|
conversation_id,
|
||||||
|
|||||||
@@ -187,30 +187,14 @@ impl ResponseStream {
|
|||||||
|
|
||||||
// Fall back to Bedrock
|
// Fall back to Bedrock
|
||||||
if *settings.bedrock_enabled.value() {
|
if *settings.bedrock_enabled.value() {
|
||||||
let auth_method = *settings.bedrock_auth_method.value();
|
|
||||||
let region = settings.bedrock_region.value().clone();
|
|
||||||
let cross_region_inference = *settings.bedrock_cross_region_inference.value();
|
|
||||||
let api_key_manager = ::ai::api_keys::ApiKeyManager::as_ref(ctx);
|
let api_key_manager = ::ai::api_keys::ApiKeyManager::as_ref(ctx);
|
||||||
let mut config = BedrockClientConfig {
|
let config = BedrockClientConfig::from_settings(settings)
|
||||||
auth_method,
|
.with_model_route(model_id, settings.bedrock_models.value())
|
||||||
profile: settings.bedrock_profile.value().clone(),
|
.with_managed_credentials(
|
||||||
region,
|
&api_key_manager.aws_credentials_refresh_strategy(),
|
||||||
access_key_id: settings.bedrock_access_key_id.value().clone(),
|
api_key_manager.aws_credentials_state(),
|
||||||
secret_access_key: settings.bedrock_secret_access_key.value().clone(),
|
);
|
||||||
session_token: None,
|
return ProviderConfig::Bedrock(config);
|
||||||
cross_region_inference,
|
|
||||||
};
|
|
||||||
|
|
||||||
if let ::ai::api_keys::AwsCredentialsState::Loaded { credentials, .. } =
|
|
||||||
api_key_manager.aws_credentials_state()
|
|
||||||
{
|
|
||||||
config.auth_method = crate::settings::BedrockAuthMethod::StaticKeys;
|
|
||||||
config.access_key_id = credentials.access_key().to_string();
|
|
||||||
config.secret_access_key = credentials.secret_key().to_string();
|
|
||||||
config.session_token = credentials.session_token().map(|s| s.to_string());
|
|
||||||
}
|
|
||||||
|
|
||||||
return ProviderConfig::Bedrock(config.with_external_fallbacks());
|
|
||||||
}
|
}
|
||||||
|
|
||||||
ProviderConfig::None
|
ProviderConfig::None
|
||||||
|
|||||||
@@ -35,8 +35,6 @@ pub struct AwsBedrockCredentialsErrorView {
|
|||||||
model_name: String,
|
model_name: String,
|
||||||
login_command: String,
|
login_command: String,
|
||||||
|
|
||||||
/// Whether auto-login was triggered, which shows a simpler "running command" message.
|
|
||||||
auto_login_triggered: bool,
|
|
||||||
|
|
||||||
// Run button
|
// Run button
|
||||||
run_button: ViewHandle<ActionButton>,
|
run_button: ViewHandle<ActionButton>,
|
||||||
@@ -52,7 +50,6 @@ impl AwsBedrockCredentialsErrorView {
|
|||||||
pub fn new(
|
pub fn new(
|
||||||
model_name: String,
|
model_name: String,
|
||||||
login_command: String,
|
login_command: String,
|
||||||
auto_login_triggered: bool,
|
|
||||||
ctx: &mut ViewContext<Self>,
|
ctx: &mut ViewContext<Self>,
|
||||||
) -> Self {
|
) -> Self {
|
||||||
// Run button
|
// Run button
|
||||||
@@ -83,7 +80,6 @@ impl AwsBedrockCredentialsErrorView {
|
|||||||
Self {
|
Self {
|
||||||
model_name,
|
model_name,
|
||||||
login_command,
|
login_command,
|
||||||
auto_login_triggered,
|
|
||||||
run_button,
|
run_button,
|
||||||
configure_button,
|
configure_button,
|
||||||
auto_login_checkbox_handle: MouseStateHandle::default(),
|
auto_login_checkbox_handle: MouseStateHandle::default(),
|
||||||
@@ -104,23 +100,6 @@ impl View for AwsBedrockCredentialsErrorView {
|
|||||||
let appearance = Appearance::as_ref(app);
|
let appearance = Appearance::as_ref(app);
|
||||||
let theme = appearance.theme();
|
let theme = appearance.theme();
|
||||||
|
|
||||||
// If auto-login was triggered, show a simple "running command" message
|
|
||||||
if self.auto_login_triggered {
|
|
||||||
return Flex::row()
|
|
||||||
.with_spacing(8.)
|
|
||||||
.with_cross_axis_alignment(CrossAxisAlignment::Center)
|
|
||||||
.with_child(
|
|
||||||
Text::new(
|
|
||||||
format!("Running `{}`...", self.login_command),
|
|
||||||
appearance.ui_font_family(),
|
|
||||||
14.,
|
|
||||||
)
|
|
||||||
.with_color(blended_colors::text_sub(theme, theme.surface_1()))
|
|
||||||
.with_selectable(false)
|
|
||||||
.finish(),
|
|
||||||
)
|
|
||||||
.finish();
|
|
||||||
}
|
|
||||||
|
|
||||||
let auto_login_enabled = *AISettings::as_ref(app).bedrock_auto_login.value();
|
let auto_login_enabled = *AISettings::as_ref(app).bedrock_auto_login.value();
|
||||||
|
|
||||||
@@ -149,11 +128,19 @@ impl View for AwsBedrockCredentialsErrorView {
|
|||||||
|
|
||||||
let make_detail_text = || {
|
let make_detail_text = || {
|
||||||
Text::new(
|
Text::new(
|
||||||
format!(
|
if self.login_command.is_empty() {
|
||||||
"Failed to authenticate with AWS Bedrock when using {}. \
|
format!(
|
||||||
Run `{}` to refresh credentials.",
|
"Failed to authenticate with AWS Bedrock when using {}. \
|
||||||
self.model_name, self.login_command
|
Refresh credentials for the AWS profile selected in Providers, then retry.",
|
||||||
),
|
self.model_name
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
format!(
|
||||||
|
"Failed to authenticate with AWS Bedrock when using {}. \
|
||||||
|
Run `{}` to refresh credentials.",
|
||||||
|
self.model_name, self.login_command
|
||||||
|
)
|
||||||
|
},
|
||||||
appearance.ui_font_family(),
|
appearance.ui_font_family(),
|
||||||
14.,
|
14.,
|
||||||
)
|
)
|
||||||
|
|||||||
+356
-224
@@ -2,12 +2,17 @@
|
|||||||
|
|
||||||
use std::io::{BufRead, BufReader, Write};
|
use std::io::{BufRead, BufReader, Write};
|
||||||
use std::net::{TcpListener, TcpStream};
|
use std::net::{TcpListener, TcpStream};
|
||||||
|
use std::sync::atomic::{AtomicBool, Ordering};
|
||||||
|
use std::sync::{Arc, Mutex, OnceLock};
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
||||||
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
|
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
|
||||||
use base64::Engine;
|
use base64::Engine;
|
||||||
use futures::channel::oneshot;
|
use futures::channel::oneshot;
|
||||||
|
use futures::StreamExt;
|
||||||
|
use galaxy_agent_rig::ChatGPTCredentialResolver;
|
||||||
use galaxyui::{Entity, ModelContext, SingletonEntity};
|
use galaxyui::{Entity, ModelContext, SingletonEntity};
|
||||||
|
use galaxyui_extras::secure_storage::AppContextExt as _;
|
||||||
use instant::Instant;
|
use instant::Instant;
|
||||||
use rand::RngCore;
|
use rand::RngCore;
|
||||||
use sha2::{Digest, Sha256};
|
use sha2::{Digest, Sha256};
|
||||||
@@ -21,6 +26,10 @@ const CALLBACK_HOST: &str = "localhost";
|
|||||||
const CALLBACK_PORTS: [u16; 2] = [1455, 1457];
|
const CALLBACK_PORTS: [u16; 2] = [1455, 1457];
|
||||||
const CALLBACK_PATH: &str = "/auth/callback";
|
const CALLBACK_PATH: &str = "/auth/callback";
|
||||||
const CALLBACK_TIMEOUT: Duration = Duration::from_secs(2 * 60);
|
const CALLBACK_TIMEOUT: Duration = Duration::from_secs(2 * 60);
|
||||||
|
const AUTH_STORAGE_KEY: &str = "galaxy.chatgpt.oauth";
|
||||||
|
const SIGN_IN_REQUIRED: &str = "ChatGPT sign-in required. Connect or reconnect ChatGPT in Providers.";
|
||||||
|
type Reply = oneshot::Sender<Result<ChatGPTAuthCredentials, String>>;
|
||||||
|
static RESOLVER: OnceLock<Mutex<Option<async_channel::Sender<Reply>>>> = OnceLock::new();
|
||||||
|
|
||||||
/// Current state of the local ChatGPT subscription connection.
|
/// Current state of the local ChatGPT subscription connection.
|
||||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||||
@@ -45,26 +54,42 @@ struct OAuthCallback {
|
|||||||
/// Coordinates browser-based OAuth authorization for ChatGPT subscriptions.
|
/// Coordinates browser-based OAuth authorization for ChatGPT subscriptions.
|
||||||
pub(crate) struct ChatGPTAuthModel {
|
pub(crate) struct ChatGPTAuthModel {
|
||||||
state: ChatGPTAuthState,
|
state: ChatGPTAuthState,
|
||||||
|
record: Option<AuthRecord>,
|
||||||
|
refresh: RefreshCoordinator,
|
||||||
|
sign_in_generation: u64,
|
||||||
|
sign_in_cancel: Option<Arc<AtomicBool>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl ChatGPTAuthModel {
|
impl ChatGPTAuthModel {
|
||||||
pub(crate) fn new() -> Self {
|
pub(crate) fn new(ctx: &mut ModelContext<Self>) -> Self {
|
||||||
let state = match has_or_import_auth_credentials() {
|
let (tx, rx) = async_channel::bounded(64);
|
||||||
Ok(()) => ChatGPTAuthState::Connected,
|
*RESOLVER.get_or_init(|| Mutex::new(None)).lock().unwrap() = Some(tx);
|
||||||
Err(error) => {
|
ctx.spawn_stream_local(rx, |model, reply, ctx| model.resolve(reply, ctx), |_, _| {});
|
||||||
log::debug!(
|
let record = ctx.secure_storage().read_value(AUTH_STORAGE_KEY).ok()
|
||||||
"[chatgpt/auth] No usable persisted ChatGPT credentials at startup: {error}"
|
.and_then(|json| serde_json::from_str::<AuthRecord>(&json).ok())
|
||||||
);
|
.filter(|record| record.credentials(now()).is_ok() || record.can_refresh());
|
||||||
ChatGPTAuthState::NotConnected
|
let state = if record.is_some() {
|
||||||
}
|
ChatGPTAuthState::Connected
|
||||||
|
} else {
|
||||||
|
ChatGPTAuthState::NotConnected
|
||||||
};
|
};
|
||||||
Self { state }
|
Self {
|
||||||
|
state,
|
||||||
|
record,
|
||||||
|
refresh: RefreshCoordinator::default(),
|
||||||
|
sign_in_generation: 0,
|
||||||
|
sign_in_cancel: None,
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
pub(crate) fn new_for_test() -> Self {
|
pub(crate) fn new_for_test() -> Self {
|
||||||
Self {
|
Self {
|
||||||
state: ChatGPTAuthState::NotConnected,
|
state: ChatGPTAuthState::NotConnected,
|
||||||
|
record: None,
|
||||||
|
refresh: RefreshCoordinator::default(),
|
||||||
|
sign_in_generation: 0,
|
||||||
|
sign_in_cancel: None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -72,22 +97,26 @@ impl ChatGPTAuthModel {
|
|||||||
&self.state
|
&self.state
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Reuses persisted credentials when possible, otherwise starts browser OAuth with PKCE.
|
fn begin_sign_in(&mut self) -> Option<u64> {
|
||||||
pub(crate) fn connect(&mut self, ctx: &mut ModelContext<Self>) {
|
|
||||||
if matches!(
|
if matches!(
|
||||||
self.state,
|
self.state,
|
||||||
ChatGPTAuthState::AwaitingBrowser | ChatGPTAuthState::ExchangingToken
|
ChatGPTAuthState::AwaitingBrowser | ChatGPTAuthState::ExchangingToken
|
||||||
) {
|
) {
|
||||||
return;
|
return None;
|
||||||
}
|
}
|
||||||
|
|
||||||
if has_or_import_auth_credentials().is_ok() {
|
self.refresh.invalidate();
|
||||||
self.state = ChatGPTAuthState::Connected;
|
self.sign_in_generation += 1;
|
||||||
ctx.emit(ChatGPTAuthModelEvent::StateChanged);
|
Some(self.sign_in_generation)
|
||||||
return;
|
}
|
||||||
}
|
|
||||||
|
|
||||||
let (redirect_uri, callback_rx) = match start_oauth_callback_server() {
|
/// Explicit connection always starts a fresh browser OAuth flow with PKCE.
|
||||||
|
pub(crate) fn connect(&mut self, ctx: &mut ModelContext<Self>) {
|
||||||
|
let Some(generation) = self.begin_sign_in() else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
let cancelled = Arc::new(AtomicBool::new(false));
|
||||||
|
let (redirect_uri, callback_rx) = match start_oauth_callback_server(cancelled.clone()) {
|
||||||
Ok(server) => server,
|
Ok(server) => server,
|
||||||
Err(error) => {
|
Err(error) => {
|
||||||
self.fail(error, ctx);
|
self.fail(error, ctx);
|
||||||
@@ -106,8 +135,8 @@ impl ChatGPTAuthModel {
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
self.state = ChatGPTAuthState::AwaitingBrowser;
|
self.sign_in_cancel = Some(cancelled);
|
||||||
ctx.emit(ChatGPTAuthModelEvent::StateChanged);
|
self.set_state(ChatGPTAuthState::AwaitingBrowser, ctx);
|
||||||
ctx.open_url(authorize_url.as_str());
|
ctx.open_url(authorize_url.as_str());
|
||||||
|
|
||||||
let _ = ctx.spawn(
|
let _ = ctx.spawn(
|
||||||
@@ -117,6 +146,9 @@ impl ChatGPTAuthModel {
|
|||||||
})?
|
})?
|
||||||
},
|
},
|
||||||
move |model, callback, ctx| {
|
move |model, callback, ctx| {
|
||||||
|
if model.sign_in_generation != generation {
|
||||||
|
return;
|
||||||
|
}
|
||||||
let callback = match callback {
|
let callback = match callback {
|
||||||
Ok(callback) => callback,
|
Ok(callback) => callback,
|
||||||
Err(error) => {
|
Err(error) => {
|
||||||
@@ -129,19 +161,27 @@ impl ChatGPTAuthModel {
|
|||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
|
||||||
model.state = ChatGPTAuthState::ExchangingToken;
|
model.set_state(ChatGPTAuthState::ExchangingToken, ctx);
|
||||||
ctx.emit(ChatGPTAuthModelEvent::StateChanged);
|
|
||||||
let _ = ctx.spawn(
|
let _ = ctx.spawn(
|
||||||
async move {
|
async move {
|
||||||
exchange_code_for_tokens(&callback.code, &code_verifier, &redirect_uri)
|
exchange_code_for_tokens(&callback.code, &code_verifier, &redirect_uri)
|
||||||
.await
|
.await
|
||||||
},
|
},
|
||||||
|model, result, ctx| match result {
|
move |model, result, ctx| {
|
||||||
Ok(()) => {
|
if model.sign_in_generation != generation {
|
||||||
model.state = ChatGPTAuthState::Connected;
|
return;
|
||||||
ctx.emit(ChatGPTAuthModelEvent::StateChanged);
|
}
|
||||||
|
match result.and_then(|record| {
|
||||||
|
save_auth_record(&record, ctx)?;
|
||||||
|
Ok(record)
|
||||||
|
}) {
|
||||||
|
Ok(record) => {
|
||||||
|
model.record = Some(record);
|
||||||
|
model.sign_in_cancel = None;
|
||||||
|
model.set_state(ChatGPTAuthState::Connected, ctx);
|
||||||
|
}
|
||||||
|
Err(error) => model.fail(error, ctx),
|
||||||
}
|
}
|
||||||
Err(error) => model.fail(error, ctx),
|
|
||||||
},
|
},
|
||||||
);
|
);
|
||||||
},
|
},
|
||||||
@@ -149,8 +189,77 @@ impl ChatGPTAuthModel {
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn fail(&mut self, error: String, ctx: &mut ModelContext<Self>) {
|
fn fail(&mut self, error: String, ctx: &mut ModelContext<Self>) {
|
||||||
self.state = ChatGPTAuthState::Failed(error);
|
if let Some(cancelled) = self.sign_in_cancel.take() {
|
||||||
ctx.emit(ChatGPTAuthModelEvent::StateChanged);
|
cancelled.store(true, Ordering::SeqCst);
|
||||||
|
}
|
||||||
|
self.set_state(ChatGPTAuthState::Failed(error), ctx);
|
||||||
|
}
|
||||||
|
|
||||||
|
fn set_state(&mut self, state: ChatGPTAuthState, ctx: &mut ModelContext<Self>) {
|
||||||
|
if self.state != state {
|
||||||
|
self.state = state;
|
||||||
|
ctx.emit(ChatGPTAuthModelEvent::StateChanged);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn cancel_sign_in(&mut self, ctx: &mut ModelContext<Self>) {
|
||||||
|
if let Some(cancelled) = self.sign_in_cancel.take() {
|
||||||
|
cancelled.store(true, Ordering::SeqCst);
|
||||||
|
self.sign_in_generation += 1;
|
||||||
|
self.refresh.invalidate();
|
||||||
|
self.set_state(ChatGPTAuthState::NotConnected, ctx);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn resolve(&mut self, reply: Reply, ctx: &mut ModelContext<Self>) {
|
||||||
|
if reply.is_canceled() {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
if matches!(self.state, ChatGPTAuthState::AwaitingBrowser | ChatGPTAuthState::ExchangingToken) {
|
||||||
|
let _ = reply.send(Err("ChatGPT sign-in is in progress.".to_owned()));
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let Some(record) = self.record.as_ref() else {
|
||||||
|
let _ = reply.send(Err(SIGN_IN_REQUIRED.to_owned()));
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
if let Ok(credentials) = record.credentials(now()) {
|
||||||
|
let _ = reply.send(Ok(credentials));
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
if !record.can_refresh() {
|
||||||
|
let _ = reply.send(Err(SIGN_IN_REQUIRED.to_owned()));
|
||||||
|
self.fail(SIGN_IN_REQUIRED.to_owned(), ctx);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let record = record.clone();
|
||||||
|
let Some(generation) = self.refresh.begin(reply) else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
let _ = ctx.spawn(
|
||||||
|
async move { refresh_tokens(&record).await },
|
||||||
|
move |model, result, ctx| {
|
||||||
|
let completed = model.refresh.complete(generation, result, |record| {
|
||||||
|
save_auth_record(record, ctx)
|
||||||
|
});
|
||||||
|
match completed {
|
||||||
|
None => {}
|
||||||
|
Some(Ok(record)) => {
|
||||||
|
model.record = Some(record);
|
||||||
|
model.set_state(ChatGPTAuthState::Connected, ctx);
|
||||||
|
}
|
||||||
|
Some(Err(error)) => {
|
||||||
|
if error.reauthenticate {
|
||||||
|
model.record = None;
|
||||||
|
if save_auth_record(&AuthRecord::default(), ctx).is_err() {
|
||||||
|
log::warn!("[chatgpt/auth] Could not persist reconnect-required state");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
model.fail(error.message, ctx);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},
|
||||||
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -194,6 +303,7 @@ fn random_hex_string(byte_count: usize) -> String {
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn start_oauth_callback_server(
|
fn start_oauth_callback_server(
|
||||||
|
cancelled: Arc<AtomicBool>,
|
||||||
) -> Result<(String, oneshot::Receiver<Result<OAuthCallback, String>>), String> {
|
) -> Result<(String, oneshot::Receiver<Result<OAuthCallback, String>>), String> {
|
||||||
let mut bind_errors = Vec::new();
|
let mut bind_errors = Vec::new();
|
||||||
let mut bound = None;
|
let mut bound = None;
|
||||||
@@ -222,6 +332,10 @@ fn start_oauth_callback_server(
|
|||||||
.spawn(move || {
|
.spawn(move || {
|
||||||
let deadline = Instant::now() + CALLBACK_TIMEOUT;
|
let deadline = Instant::now() + CALLBACK_TIMEOUT;
|
||||||
loop {
|
loop {
|
||||||
|
if cancelled.load(Ordering::SeqCst) {
|
||||||
|
let _ = callback_tx.send(Err("ChatGPT sign-in was cancelled.".to_owned()));
|
||||||
|
return;
|
||||||
|
}
|
||||||
match listener.accept() {
|
match listener.accept() {
|
||||||
Ok((stream, _)) => {
|
Ok((stream, _)) => {
|
||||||
let _ = callback_tx.send(handle_callback_connection(stream));
|
let _ = callback_tx.send(handle_callback_connection(stream));
|
||||||
@@ -323,26 +437,19 @@ fn parse_callback_request_line(request_line: &str) -> Result<OAuthCallback, Stri
|
|||||||
|
|
||||||
let mut code = None;
|
let mut code = None;
|
||||||
let mut state = None;
|
let mut state = None;
|
||||||
let mut oauth_error = None;
|
let mut oauth_error = false;
|
||||||
let mut error_description = None;
|
|
||||||
for (key, value) in url.query_pairs() {
|
for (key, value) in url.query_pairs() {
|
||||||
match key.as_ref() {
|
match key.as_ref() {
|
||||||
"code" if !value.is_empty() => code = Some(value.into_owned()),
|
"code" if !value.is_empty() => code = Some(value.into_owned()),
|
||||||
"state" if !value.is_empty() => state = Some(value.into_owned()),
|
"state" if !value.is_empty() => state = Some(value.into_owned()),
|
||||||
"error" if !value.is_empty() => oauth_error = Some(value.into_owned()),
|
"error" if !value.is_empty() => oauth_error = true,
|
||||||
"error_description" if !value.is_empty() => {
|
|
||||||
error_description = Some(value.into_owned())
|
|
||||||
}
|
|
||||||
"code" | "state" | "error" | "error_description" => {}
|
"code" | "state" | "error" | "error_description" => {}
|
||||||
_ => {}
|
_ => {}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
if let Some(error) = oauth_error {
|
if oauth_error {
|
||||||
return Err(format!(
|
return Err("ChatGPT browser authorization failed. Please try signing in again.".to_owned());
|
||||||
"ChatGPT authorization failed: {error} ({})",
|
|
||||||
error_description.as_deref().unwrap_or("no description")
|
|
||||||
));
|
|
||||||
}
|
}
|
||||||
Ok(OAuthCallback {
|
Ok(OAuthCallback {
|
||||||
code: code.ok_or("ChatGPT OAuth callback did not include an authorization code.")?,
|
code: code.ok_or("ChatGPT OAuth callback did not include an authorization code.")?,
|
||||||
@@ -358,157 +465,101 @@ fn callback_response_page(title: &str, message: &str, is_error: bool) -> String
|
|||||||
.replace("__ACCENT__", accent)
|
.replace("__ACCENT__", accent)
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
pub(crate) use galaxy_agent_rig::ChatGPTSubscriptionCredentials as ChatGPTAuthCredentials;
|
||||||
pub(crate) struct ChatGPTAuthCredentials {
|
|
||||||
pub(crate) access_token: String,
|
|
||||||
pub(crate) account_id: Option<String>,
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Attempts to read tokens from `~/.codex/auth.json` and write them to Rig's auth file.
|
pub(crate) fn credential_resolver() -> ChatGPTCredentialResolver {
|
||||||
/// Returns `Ok(())` if credentials were found and successfully imported.
|
Arc::new(|| Box::pin(load_or_refresh_auth_credentials()))
|
||||||
pub(crate) fn import_codex_credentials() -> Result<(), String> {
|
|
||||||
let codex_path = codex_auth_file_path().ok_or("Cannot determine codex auth path")?;
|
|
||||||
let bytes = std::fs::read(&codex_path).map_err(|e| format!("{e}"))?;
|
|
||||||
let doc: serde_json::Value = serde_json::from_slice(&bytes).map_err(|e| format!("{e}"))?;
|
|
||||||
|
|
||||||
let tokens = doc.get("tokens").ok_or("No tokens object")?;
|
|
||||||
let access_token = tokens
|
|
||||||
.get("access_token")
|
|
||||||
.and_then(|value| value.as_str())
|
|
||||||
.filter(|token| !token.is_empty());
|
|
||||||
let refresh_token = tokens
|
|
||||||
.get("refresh_token")
|
|
||||||
.and_then(|value| value.as_str())
|
|
||||||
.filter(|token| !token.is_empty())
|
|
||||||
.map(ToOwned::to_owned);
|
|
||||||
if access_token.is_none() && refresh_token.is_none() {
|
|
||||||
return Err("Codex auth file does not contain an access or refresh token".to_string());
|
|
||||||
}
|
|
||||||
let expires_at = access_token.and_then(extract_expiration_timestamp);
|
|
||||||
let id_token = tokens
|
|
||||||
.get("id_token")
|
|
||||||
.and_then(|v| v.as_str())
|
|
||||||
.map(ToOwned::to_owned);
|
|
||||||
let account_id = tokens
|
|
||||||
.get("account_id")
|
|
||||||
.and_then(|v| v.as_str())
|
|
||||||
.map(ToOwned::to_owned)
|
|
||||||
.or_else(|| extract_account_id(id_token.as_deref()))
|
|
||||||
.or_else(|| extract_account_id(access_token));
|
|
||||||
|
|
||||||
let record = AuthRecord {
|
|
||||||
access_token: access_token.map(ToOwned::to_owned),
|
|
||||||
refresh_token,
|
|
||||||
id_token,
|
|
||||||
expires_at,
|
|
||||||
account_id,
|
|
||||||
};
|
|
||||||
|
|
||||||
write_auth_file(&record)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn load_or_refresh_auth_credentials() -> Result<ChatGPTAuthCredentials, String> {
|
pub(crate) async fn load_or_refresh_auth_credentials() -> Result<ChatGPTAuthCredentials, String> {
|
||||||
if let Ok(credentials) = load_auth_credentials() {
|
let (tx, rx) = oneshot::channel();
|
||||||
return Ok(credentials);
|
RESOLVER.get_or_init(|| Mutex::new(None)).lock().unwrap()
|
||||||
}
|
.as_ref().ok_or("ChatGPT authorization service is unavailable.")?
|
||||||
|
.try_send(tx).map_err(|_| "ChatGPT authorization service is busy or unavailable.")?;
|
||||||
let _ = import_codex_credentials();
|
match futures::future::select(
|
||||||
if let Ok(credentials) = load_auth_credentials() {
|
Box::pin(rx),
|
||||||
return Ok(credentials);
|
Box::pin(async_io::Timer::after(Duration::from_secs(35))),
|
||||||
}
|
).await {
|
||||||
|
futures::future::Either::Left((result, _)) => {
|
||||||
galaxy_agent_rig::refresh_chatgpt_subscription_credentials().await?;
|
result.map_err(|_| "ChatGPT authorization service stopped.".to_owned())?
|
||||||
load_auth_credentials()
|
}
|
||||||
}
|
futures::future::Either::Right(_) => {
|
||||||
|
Err("ChatGPT credential refresh timed out. Retry or sign in again.".to_owned())
|
||||||
fn has_or_import_auth_credentials() -> Result<(), String> {
|
}
|
||||||
if has_auth_credentials() {
|
|
||||||
return Ok(());
|
|
||||||
}
|
|
||||||
import_codex_credentials()?;
|
|
||||||
if has_auth_credentials() {
|
|
||||||
Ok(())
|
|
||||||
} else {
|
|
||||||
Err("ChatGPT credentials are missing an access or refresh token.".to_string())
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn has_auth_credentials() -> bool {
|
#[derive(Clone, Debug)]
|
||||||
let Some(path) = auth_file_path() else {
|
struct AuthFailure {
|
||||||
return false;
|
message: String,
|
||||||
};
|
reauthenticate: bool,
|
||||||
let Ok(bytes) = std::fs::read(path) else {
|
|
||||||
return false;
|
|
||||||
};
|
|
||||||
let Ok(record) = serde_json::from_slice::<AuthRecord>(&bytes) else {
|
|
||||||
return false;
|
|
||||||
};
|
|
||||||
record
|
|
||||||
.refresh_token
|
|
||||||
.as_deref()
|
|
||||||
.is_some_and(|token| !token.trim().is_empty())
|
|
||||||
|| record.access_token.as_deref().is_some_and(|token| {
|
|
||||||
!token.trim().is_empty()
|
|
||||||
&& !token_is_expired(
|
|
||||||
record
|
|
||||||
.expires_at
|
|
||||||
.or_else(|| extract_expiration_timestamp(token)),
|
|
||||||
)
|
|
||||||
})
|
|
||||||
}
|
}
|
||||||
|
|
||||||
fn load_auth_credentials() -> Result<ChatGPTAuthCredentials, String> {
|
#[derive(Default)]
|
||||||
let path = auth_file_path().ok_or("Cannot determine ChatGPT auth file path")?;
|
struct RefreshCoordinator {
|
||||||
let bytes = std::fs::read(&path)
|
generation: u64,
|
||||||
.map_err(|error| format!("Failed to read {}: {error}", path.display()))?;
|
active: Option<(u64, Vec<Reply>)>,
|
||||||
let record: AuthRecord = serde_json::from_slice(&bytes)
|
}
|
||||||
.map_err(|error| format!("Failed to parse {}: {error}", path.display()))?;
|
|
||||||
let access_token = record
|
|
||||||
.access_token
|
|
||||||
.as_deref()
|
|
||||||
.filter(|token| !token.trim().is_empty())
|
|
||||||
.ok_or("ChatGPT auth file does not contain an access token")?;
|
|
||||||
|
|
||||||
let expires_at = record
|
impl RefreshCoordinator {
|
||||||
.expires_at
|
fn begin(&mut self, reply: Reply) -> Option<u64> {
|
||||||
.or_else(|| extract_expiration_timestamp(access_token));
|
if let Some((_, replies)) = &mut self.active {
|
||||||
if token_is_expired(expires_at) {
|
replies.push(reply);
|
||||||
return Err("ChatGPT access token is expired".to_string());
|
return None;
|
||||||
|
}
|
||||||
|
self.generation += 1;
|
||||||
|
self.active = Some((self.generation, vec![reply]));
|
||||||
|
Some(self.generation)
|
||||||
}
|
}
|
||||||
|
|
||||||
let account_id = record
|
fn invalidate(&mut self) {
|
||||||
.account_id
|
self.generation += 1;
|
||||||
.clone()
|
if let Some((_, replies)) = self.active.take() {
|
||||||
.filter(|account_id| !account_id.trim().is_empty())
|
for reply in replies {
|
||||||
.or_else(|| extract_account_id(record.id_token.as_deref()))
|
let _ = reply.send(Err("ChatGPT session was replaced. Please sign in again.".to_owned()));
|
||||||
.or_else(|| extract_account_id(Some(access_token)));
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
Ok(ChatGPTAuthCredentials {
|
fn complete(
|
||||||
access_token: access_token.to_string(),
|
&mut self,
|
||||||
account_id,
|
generation: u64,
|
||||||
})
|
result: Result<AuthRecord, AuthFailure>,
|
||||||
}
|
persist: impl FnOnce(&AuthRecord) -> Result<(), String>,
|
||||||
|
) -> Option<Result<AuthRecord, AuthFailure>> {
|
||||||
fn codex_auth_file_path() -> Option<std::path::PathBuf> {
|
if self.active.as_ref().is_none_or(|(active, _)| *active != generation) {
|
||||||
if let Some(codex_home) = std::env::var_os("CODEX_HOME") {
|
return None;
|
||||||
return Some(std::path::PathBuf::from(codex_home).join("auth.json"));
|
}
|
||||||
|
let (_, replies) = self.active.take()?;
|
||||||
|
let result = result.and_then(|record| {
|
||||||
|
let credentials = record.access_context().map_err(|message| AuthFailure {
|
||||||
|
message, reauthenticate: true,
|
||||||
|
})?;
|
||||||
|
persist(&record).map_err(|_| AuthFailure {
|
||||||
|
message: "Could not securely save refreshed ChatGPT credentials. Please sign in again.".to_owned(),
|
||||||
|
reauthenticate: true,
|
||||||
|
})?;
|
||||||
|
Ok((record, credentials))
|
||||||
|
});
|
||||||
|
for reply in replies {
|
||||||
|
let response = result.as_ref().map(|(_, credentials)| credentials.clone())
|
||||||
|
.map_err(|error| error.message.clone());
|
||||||
|
let _ = reply.send(response);
|
||||||
|
}
|
||||||
|
Some(result.map(|(record, _)| record))
|
||||||
}
|
}
|
||||||
std::env::var_os("HOME").map(|h| std::path::PathBuf::from(h).join(".codex").join("auth.json"))
|
|
||||||
}
|
}
|
||||||
|
|
||||||
fn token_is_expired(expires_at: Option<i64>) -> bool {
|
fn now() -> i64 {
|
||||||
let now = std::time::SystemTime::now()
|
std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH)
|
||||||
.duration_since(std::time::UNIX_EPOCH)
|
.map(|duration| duration.as_secs() as i64).unwrap_or_default()
|
||||||
.map(|duration| duration.as_secs() as i64)
|
|
||||||
.unwrap_or_default();
|
|
||||||
expires_at.is_none_or(|expires_at| now >= expires_at - 60)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn exchange_code_for_tokens(
|
async fn exchange_code_for_tokens(
|
||||||
code: &str,
|
code: &str,
|
||||||
code_verifier: &str,
|
code_verifier: &str,
|
||||||
redirect_uri: &str,
|
redirect_uri: &str,
|
||||||
) -> Result<(), String> {
|
) -> Result<AuthRecord, String> {
|
||||||
let body = url::form_urlencoded::Serializer::new(String::new())
|
let body = url::form_urlencoded::Serializer::new(String::new())
|
||||||
.append_pair("grant_type", "authorization_code")
|
.append_pair("grant_type", "authorization_code")
|
||||||
.append_pair("client_id", CHATGPT_CLIENT_ID)
|
.append_pair("client_id", CHATGPT_CLIENT_ID)
|
||||||
@@ -516,7 +567,34 @@ async fn exchange_code_for_tokens(
|
|||||||
.append_pair("redirect_uri", redirect_uri)
|
.append_pair("redirect_uri", redirect_uri)
|
||||||
.append_pair("code_verifier", code_verifier)
|
.append_pair("code_verifier", code_verifier)
|
||||||
.finish();
|
.finish();
|
||||||
let response = reqwest::Client::new()
|
let tokens = token_request(body).await.map_err(|error| error.message)?;
|
||||||
|
record_from_tokens(tokens, None, now()).map_err(|error| error.message)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn refresh_tokens(previous: &AuthRecord) -> Result<AuthRecord, AuthFailure> {
|
||||||
|
let refresh_token = previous.refresh_token.as_deref().filter(|token| !token.trim().is_empty())
|
||||||
|
.ok_or_else(|| AuthFailure {
|
||||||
|
message: SIGN_IN_REQUIRED.to_owned(),
|
||||||
|
reauthenticate: true,
|
||||||
|
})?;
|
||||||
|
let body = url::form_urlencoded::Serializer::new(String::new())
|
||||||
|
.append_pair("grant_type", "refresh_token")
|
||||||
|
.append_pair("client_id", CHATGPT_CLIENT_ID)
|
||||||
|
.append_pair("refresh_token", refresh_token)
|
||||||
|
.finish();
|
||||||
|
let tokens = token_request(body).await?;
|
||||||
|
record_from_tokens(tokens, Some(previous), now())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn token_request(body: String) -> Result<TokenResponse, AuthFailure> {
|
||||||
|
let client = reqwest::Client::builder()
|
||||||
|
.timeout(Duration::from_secs(30))
|
||||||
|
.redirect(reqwest::redirect::Policy::none())
|
||||||
|
.build().map_err(|_| AuthFailure {
|
||||||
|
message: "Could not create the ChatGPT authorization client.".to_owned(),
|
||||||
|
reauthenticate: false,
|
||||||
|
})?;
|
||||||
|
let response = client
|
||||||
.post(CHATGPT_TOKEN_URL)
|
.post(CHATGPT_TOKEN_URL)
|
||||||
.header(
|
.header(
|
||||||
reqwest::header::CONTENT_TYPE,
|
reqwest::header::CONTENT_TYPE,
|
||||||
@@ -525,67 +603,83 @@ async fn exchange_code_for_tokens(
|
|||||||
.body(body)
|
.body(body)
|
||||||
.send()
|
.send()
|
||||||
.await
|
.await
|
||||||
.map_err(|error| format!("ChatGPT token exchange request failed: {error}"))?;
|
.map_err(|_| AuthFailure {
|
||||||
|
message: "ChatGPT authorization request failed. Retry the connection.".to_owned(),
|
||||||
let status = response.status();
|
reauthenticate: false,
|
||||||
if !status.is_success() {
|
})?;
|
||||||
let body = response.text().await.unwrap_or_default();
|
let status = response.status().as_u16();
|
||||||
return Err(format!("ChatGPT token exchange failed: {status} {body}"));
|
let mut bytes = Vec::new();
|
||||||
|
let mut stream = response.bytes_stream();
|
||||||
|
while let Some(chunk) = stream.next().await {
|
||||||
|
let chunk = chunk.map_err(|_| AuthFailure {
|
||||||
|
message: "Could not read the ChatGPT authorization response.".to_owned(),
|
||||||
|
reauthenticate: (200..300).contains(&status),
|
||||||
|
})?;
|
||||||
|
if bytes.len().saturating_add(chunk.len()) > 64 * 1024 {
|
||||||
|
return Err(AuthFailure {
|
||||||
|
message: "ChatGPT authorization returned an oversized response.".to_owned(),
|
||||||
|
reauthenticate: (200..300).contains(&status),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
bytes.extend_from_slice(&chunk);
|
||||||
}
|
}
|
||||||
let tokens: TokenResponse = response
|
parse_token_response(status, &bytes)
|
||||||
.json()
|
}
|
||||||
.await
|
|
||||||
.map_err(|error| format!("Could not parse the ChatGPT token response: {error}"))?;
|
fn parse_token_response(status: u16, bytes: &[u8]) -> Result<TokenResponse, AuthFailure> {
|
||||||
|
if (200..300).contains(&status) {
|
||||||
|
return serde_json::from_slice(bytes).map_err(|_| AuthFailure {
|
||||||
|
message: "ChatGPT returned invalid credentials. Please sign in again.".to_owned(),
|
||||||
|
reauthenticate: true,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
let body = serde_json::from_slice::<serde_json::Value>(bytes).unwrap_or_default();
|
||||||
|
let code = body.get("error").and_then(|error| {
|
||||||
|
error.as_str().or_else(|| error.get("code").and_then(|code| code.as_str()))
|
||||||
|
});
|
||||||
|
let reauthenticate = matches!(status, 401 | 403)
|
||||||
|
|| (status == 400 && matches!(code,
|
||||||
|
Some("invalid_grant" | "refresh_token_reused" | "refresh_token_expired" | "refresh_token_invalid")
|
||||||
|
));
|
||||||
|
let message = if reauthenticate {
|
||||||
|
format!("ChatGPT authorization expired or was rejected (HTTP {status}). Please sign in again.")
|
||||||
|
} else {
|
||||||
|
format!("ChatGPT authorization failed (HTTP {status}). Retry the connection.")
|
||||||
|
};
|
||||||
|
Err(AuthFailure { message, reauthenticate })
|
||||||
|
}
|
||||||
|
|
||||||
|
fn record_from_tokens(
|
||||||
|
tokens: TokenResponse,
|
||||||
|
previous: Option<&AuthRecord>,
|
||||||
|
now: i64,
|
||||||
|
) -> Result<AuthRecord, AuthFailure> {
|
||||||
let expires_at = extract_expiration_timestamp(&tokens.access_token).or_else(|| {
|
let expires_at = extract_expiration_timestamp(&tokens.access_token).or_else(|| {
|
||||||
tokens.expires_in.map(|expires_in| {
|
tokens.expires_in.and_then(|expires| i64::try_from(expires).ok())
|
||||||
std::time::SystemTime::now()
|
.map(|expires| now.saturating_add(expires))
|
||||||
.duration_since(std::time::UNIX_EPOCH)
|
|
||||||
.map(|duration| duration.as_secs() as i64)
|
|
||||||
.unwrap_or_default()
|
|
||||||
+ expires_in as i64
|
|
||||||
})
|
|
||||||
});
|
});
|
||||||
let account_id = extract_account_id(tokens.id_token.as_deref())
|
let account_id = extract_account_id(tokens.id_token.as_deref())
|
||||||
.or_else(|| extract_account_id(Some(&tokens.access_token)));
|
.or_else(|| extract_account_id(Some(&tokens.access_token)))
|
||||||
write_auth_file(&AuthRecord {
|
.or_else(|| previous.and_then(|record| record.account_id.clone()));
|
||||||
|
let record = AuthRecord {
|
||||||
access_token: Some(tokens.access_token),
|
access_token: Some(tokens.access_token),
|
||||||
refresh_token: tokens.refresh_token,
|
refresh_token: tokens.refresh_token.or_else(|| previous.and_then(|record| record.refresh_token.clone())),
|
||||||
id_token: tokens.id_token,
|
id_token: tokens.id_token.or_else(|| previous.and_then(|record| record.id_token.clone())),
|
||||||
expires_at,
|
expires_at,
|
||||||
account_id,
|
account_id,
|
||||||
})
|
};
|
||||||
|
record.credentials(now).map_err(|_| AuthFailure {
|
||||||
|
message: "ChatGPT returned invalid or expired credentials. Please sign in again.".to_owned(),
|
||||||
|
reauthenticate: true,
|
||||||
|
})?;
|
||||||
|
Ok(record)
|
||||||
}
|
}
|
||||||
|
|
||||||
fn write_auth_file(record: &AuthRecord) -> Result<(), String> {
|
fn save_auth_record(record: &AuthRecord, ctx: &galaxyui::AppContext) -> Result<(), String> {
|
||||||
let path = auth_file_path().ok_or("Cannot determine auth file path")?;
|
let json = serde_json::to_string(record)
|
||||||
if let Some(parent) = path.parent() {
|
.map_err(|_| "Could not encode ChatGPT credentials.".to_owned())?;
|
||||||
std::fs::create_dir_all(parent)
|
ctx.secure_storage().write_value_with_owner_only_fallback(AUTH_STORAGE_KEY, &json)
|
||||||
.map_err(|e| format!("Failed to create auth directory: {e}"))?;
|
.map_err(|_| "Could not securely save ChatGPT credentials. Please sign in again.".to_owned())
|
||||||
}
|
|
||||||
let json =
|
|
||||||
serde_json::to_vec_pretty(record).map_err(|e| format!("Failed to serialize auth: {e}"))?;
|
|
||||||
std::fs::write(&path, json).map_err(|e| format!("Failed to write auth file: {e}"))?;
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
|
|
||||||
fn auth_file_path() -> Option<std::path::PathBuf> {
|
|
||||||
#[cfg(target_os = "windows")]
|
|
||||||
{
|
|
||||||
std::env::var_os("APPDATA").map(|d| {
|
|
||||||
std::path::PathBuf::from(d)
|
|
||||||
.join("chatgpt")
|
|
||||||
.join("auth.json")
|
|
||||||
})
|
|
||||||
}
|
|
||||||
#[cfg(not(target_os = "windows"))]
|
|
||||||
{
|
|
||||||
std::env::var_os("XDG_CONFIG_HOME")
|
|
||||||
.map(std::path::PathBuf::from)
|
|
||||||
.or_else(|| {
|
|
||||||
std::env::var_os("HOME").map(|h| std::path::PathBuf::from(h).join(".config"))
|
|
||||||
})
|
|
||||||
.map(|d| d.join("chatgpt").join("auth.json"))
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
fn extract_expiration_timestamp(token: &str) -> Option<i64> {
|
fn extract_expiration_timestamp(token: &str) -> Option<i64> {
|
||||||
@@ -633,7 +727,7 @@ struct TokenResponse {
|
|||||||
expires_in: Option<u64>,
|
expires_in: Option<u64>,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(serde::Deserialize, serde::Serialize)]
|
#[derive(Clone, Default, serde::Deserialize, serde::Serialize)]
|
||||||
struct AuthRecord {
|
struct AuthRecord {
|
||||||
access_token: Option<String>,
|
access_token: Option<String>,
|
||||||
refresh_token: Option<String>,
|
refresh_token: Option<String>,
|
||||||
@@ -641,3 +735,41 @@ struct AuthRecord {
|
|||||||
expires_at: Option<i64>,
|
expires_at: Option<i64>,
|
||||||
account_id: Option<String>,
|
account_id: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
impl std::fmt::Debug for AuthRecord {
|
||||||
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||||
|
f.debug_struct("ChatGPTAuthRecord")
|
||||||
|
.field("credentials", &"<redacted>")
|
||||||
|
.field("expires_at", &self.expires_at)
|
||||||
|
.finish()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl AuthRecord {
|
||||||
|
fn can_refresh(&self) -> bool {
|
||||||
|
self.refresh_token.as_deref().is_some_and(|token| !token.trim().is_empty())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn access_context(&self) -> Result<ChatGPTAuthCredentials, String> {
|
||||||
|
let access_token = self.access_token.as_deref().filter(|token| !token.trim().is_empty())
|
||||||
|
.ok_or_else(|| SIGN_IN_REQUIRED.to_owned())?;
|
||||||
|
let account_id = self.account_id.clone().filter(|id| !id.trim().is_empty())
|
||||||
|
.or_else(|| extract_account_id(self.id_token.as_deref()))
|
||||||
|
.or_else(|| extract_account_id(Some(access_token)));
|
||||||
|
Ok(ChatGPTAuthCredentials { access_token: access_token.to_owned(), account_id })
|
||||||
|
}
|
||||||
|
|
||||||
|
fn credentials(&self, now: i64) -> Result<ChatGPTAuthCredentials, String> {
|
||||||
|
let expiry = self.expires_at.or_else(|| {
|
||||||
|
self.access_token.as_deref().and_then(extract_expiration_timestamp)
|
||||||
|
});
|
||||||
|
if expiry.is_none_or(|expiry| now >= expiry.saturating_sub(60)) {
|
||||||
|
return Err(SIGN_IN_REQUIRED.to_owned());
|
||||||
|
}
|
||||||
|
self.access_context()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
#[path = "chatgpt_auth_tests.rs"]
|
||||||
|
mod tests;
|
||||||
|
|||||||
@@ -0,0 +1,171 @@
|
|||||||
|
use std::cell::Cell;
|
||||||
|
|
||||||
|
use futures::FutureExt;
|
||||||
|
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
fn record(access: &str, refresh: &str) -> AuthRecord {
|
||||||
|
AuthRecord {
|
||||||
|
access_token: Some(access.to_owned()),
|
||||||
|
refresh_token: Some(refresh.to_owned()),
|
||||||
|
id_token: None,
|
||||||
|
expires_at: Some(10_000),
|
||||||
|
account_id: Some("test-account".to_owned()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn concurrent_requests_share_one_refresh_and_wait_for_persistence() {
|
||||||
|
let mut coordinator = RefreshCoordinator::default();
|
||||||
|
let (first_tx, mut first_rx) = oneshot::channel();
|
||||||
|
let (second_tx, mut second_rx) = oneshot::channel();
|
||||||
|
let generation = coordinator.begin(first_tx).unwrap();
|
||||||
|
assert_eq!(coordinator.begin(second_tx), None);
|
||||||
|
let saved = Cell::new(false);
|
||||||
|
let result = coordinator.complete(generation, Ok(record("rotated-access", "rotated-refresh")), |record| {
|
||||||
|
assert_eq!(first_rx.try_recv().unwrap(), None);
|
||||||
|
assert_eq!(second_rx.try_recv().unwrap(), None);
|
||||||
|
assert_eq!(record.refresh_token.as_deref(), Some("rotated-refresh"));
|
||||||
|
saved.set(true);
|
||||||
|
Ok(())
|
||||||
|
});
|
||||||
|
assert!(result.unwrap().is_ok());
|
||||||
|
assert!(saved.get());
|
||||||
|
let first = first_rx.now_or_never().unwrap().unwrap().unwrap();
|
||||||
|
let second = second_rx.now_or_never().unwrap().unwrap().unwrap();
|
||||||
|
assert_eq!(first, second);
|
||||||
|
assert_eq!(first.access_token, "rotated-access");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn failed_persistence_never_publishes_rotated_credentials() {
|
||||||
|
let mut coordinator = RefreshCoordinator::default();
|
||||||
|
let (tx, rx) = oneshot::channel();
|
||||||
|
let generation = coordinator.begin(tx).unwrap();
|
||||||
|
let result = coordinator.complete(generation, Ok(record("rotated", "new-refresh")), |_| {
|
||||||
|
Err("storage unavailable".to_owned())
|
||||||
|
}).unwrap().unwrap_err();
|
||||||
|
assert!(result.reauthenticate);
|
||||||
|
assert!(rx.now_or_never().unwrap().unwrap().is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn replaced_sign_in_rejects_late_refresh_without_overwriting_new_session() {
|
||||||
|
let mut coordinator = RefreshCoordinator::default();
|
||||||
|
let (old_tx, old_rx) = oneshot::channel();
|
||||||
|
let old = coordinator.begin(old_tx).unwrap();
|
||||||
|
coordinator.invalidate();
|
||||||
|
assert!(old_rx.now_or_never().unwrap().unwrap().is_err());
|
||||||
|
let (new_tx, new_rx) = oneshot::channel();
|
||||||
|
let current = coordinator.begin(new_tx).unwrap();
|
||||||
|
assert_ne!(old, current);
|
||||||
|
assert!(coordinator.complete(old, Ok(record("stale", "stale-refresh")), |_| {
|
||||||
|
panic!("A stale refresh must not write credentials")
|
||||||
|
}).is_none());
|
||||||
|
assert!(coordinator.complete(current, Ok(record("current", "current-refresh")), |_| Ok(()))
|
||||||
|
.unwrap().is_ok());
|
||||||
|
assert_eq!(new_rx.now_or_never().unwrap().unwrap().unwrap().access_token, "current");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn reconnect_starts_a_new_generation_even_with_a_cached_refresh_token() {
|
||||||
|
let mut model = ChatGPTAuthModel::new_for_test();
|
||||||
|
model.record = Some(record("old-access", "already-used-refresh"));
|
||||||
|
model.state = ChatGPTAuthState::Connected;
|
||||||
|
let first = model.begin_sign_in().unwrap();
|
||||||
|
model.state = ChatGPTAuthState::AwaitingBrowser;
|
||||||
|
assert_eq!(model.begin_sign_in(), None);
|
||||||
|
model.state = ChatGPTAuthState::Failed("refresh_token_reused".to_owned());
|
||||||
|
let replacement = model.begin_sign_in().unwrap();
|
||||||
|
assert!(replacement > first);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn rotated_tokens_replace_previous_tokens_and_omitted_refresh_tokens_are_preserved() {
|
||||||
|
let previous = record("old-access", "old-refresh");
|
||||||
|
for refresh in [None, Some("new-refresh".to_owned())] {
|
||||||
|
let result = record_from_tokens(TokenResponse {
|
||||||
|
access_token: "new-access".to_owned(),
|
||||||
|
refresh_token: refresh.clone(),
|
||||||
|
id_token: None,
|
||||||
|
expires_in: Some(3600),
|
||||||
|
}, Some(&previous), 100).unwrap();
|
||||||
|
assert_eq!(result.access_token.as_deref(), Some("new-access"));
|
||||||
|
assert_eq!(result.refresh_token, refresh.or_else(|| Some("old-refresh".to_owned())));
|
||||||
|
assert_eq!(result.expires_at, Some(3700));
|
||||||
|
assert_eq!(result.account_id.as_deref(), Some("test-account"));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn rejected_refresh_tokens_require_reauthentication_without_echoing_tokens() {
|
||||||
|
for code in ["refresh_token_reused", "invalid_grant", "refresh_token_expired"] {
|
||||||
|
let body = serde_json::to_vec(&serde_json::json!({
|
||||||
|
"error": {"code": code, "message": "sensitive-refresh-token"}
|
||||||
|
})).unwrap();
|
||||||
|
let error = parse_token_response(401, &body).err().unwrap();
|
||||||
|
assert!(error.reauthenticate);
|
||||||
|
assert!(error.message.contains("sign in"));
|
||||||
|
assert!(!error.message.contains("sensitive-refresh-token"));
|
||||||
|
}
|
||||||
|
let error = parse_token_response(400, br#"{"error":"invalid_grant"}"#).err().unwrap();
|
||||||
|
assert!(error.reauthenticate);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn transient_refresh_failure_preserves_the_session_and_does_not_save_rejected_data() {
|
||||||
|
let error = parse_token_response(503, br#"{"message":"sensitive-data"}"#).err().unwrap();
|
||||||
|
assert!(!error.reauthenticate);
|
||||||
|
assert!(!error.message.contains("sensitive-data"));
|
||||||
|
let mut coordinator = RefreshCoordinator::default();
|
||||||
|
let (tx, rx) = oneshot::channel();
|
||||||
|
let generation = coordinator.begin(tx).unwrap();
|
||||||
|
let result = coordinator.complete(generation, Err(error), |_| {
|
||||||
|
panic!("Transient failures must not overwrite stored credentials")
|
||||||
|
}).unwrap().unwrap_err();
|
||||||
|
assert!(!result.reauthenticate);
|
||||||
|
assert!(rx.now_or_never().unwrap().unwrap().is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn credential_debug_output_is_redacted_and_expiry_has_a_refresh_margin() {
|
||||||
|
let credentials = ChatGPTAuthCredentials {
|
||||||
|
access_token: "sensitive-access-token".to_owned(),
|
||||||
|
account_id: Some("account".to_owned()),
|
||||||
|
};
|
||||||
|
assert!(!format!("{credentials:?}").contains("sensitive-access-token"));
|
||||||
|
let credentials = record("sensitive-access-token", "sensitive-refresh-token");
|
||||||
|
assert!(!format!("{credentials:?}").contains("sensitive-access-token"));
|
||||||
|
assert!(!format!("{credentials:?}").contains("sensitive-refresh-token"));
|
||||||
|
assert!(credentials.credentials(100).is_ok());
|
||||||
|
assert!(credentials.credentials(9940).is_err());
|
||||||
|
assert!(AuthRecord::default().credentials(100).is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn invalid_token_payloads_are_not_saved_or_exposed() {
|
||||||
|
assert!(parse_token_response(200, b"invalid-sensitive-data").is_err());
|
||||||
|
assert!(record_from_tokens(TokenResponse {
|
||||||
|
access_token: String::new(),
|
||||||
|
refresh_token: Some("refresh".to_owned()),
|
||||||
|
id_token: None,
|
||||||
|
expires_in: Some(3600),
|
||||||
|
}, None, 100).is_err());
|
||||||
|
assert!(record_from_tokens(TokenResponse {
|
||||||
|
access_token: "opaque-token-without-expiry".to_owned(),
|
||||||
|
refresh_token: None,
|
||||||
|
id_token: None,
|
||||||
|
expires_in: None,
|
||||||
|
}, None, 100).is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn browser_authorization_keeps_pkce_and_state_validation() {
|
||||||
|
let url = build_authorize_url("http://localhost:1455/auth/callback", "challenge", "state").unwrap();
|
||||||
|
let params = url.query_pairs().collect::<std::collections::HashMap<_, _>>();
|
||||||
|
assert_eq!(params.get("code_challenge_method").map(|v| v.as_ref()), Some("S256"));
|
||||||
|
assert_eq!(params.get("state").map(|v| v.as_ref()), Some("state"));
|
||||||
|
let callback = parse_callback_request_line("GET /auth/callback?code=code&state=state HTTP/1.1").unwrap();
|
||||||
|
assert_eq!(callback.state, "state");
|
||||||
|
assert!(parse_callback_request_line("GET /auth/callback?code=code HTTP/1.1").is_err());
|
||||||
|
}
|
||||||
@@ -181,28 +181,14 @@ impl CrosscheckReviewer {
|
|||||||
|
|
||||||
// Fall back to Bedrock via external config
|
// Fall back to Bedrock via external config
|
||||||
if *settings.bedrock_enabled.value() {
|
if *settings.bedrock_enabled.value() {
|
||||||
let auth_method = *settings.bedrock_auth_method.value();
|
|
||||||
let api_key_manager = ::ai::api_keys::ApiKeyManager::as_ref(ctx);
|
let api_key_manager = ::ai::api_keys::ApiKeyManager::as_ref(ctx);
|
||||||
let mut config = crate::ai::provider::client::BedrockClientConfig {
|
let config = crate::ai::provider::client::BedrockClientConfig::from_settings(settings)
|
||||||
auth_method,
|
.with_model_route(model_id, settings.bedrock_models.value())
|
||||||
profile: settings.bedrock_profile.value().clone(),
|
.with_managed_credentials(
|
||||||
region: settings.bedrock_region.value().clone(),
|
&api_key_manager.aws_credentials_refresh_strategy(),
|
||||||
access_key_id: settings.bedrock_access_key_id.value().clone(),
|
api_key_manager.aws_credentials_state(),
|
||||||
secret_access_key: settings.bedrock_secret_access_key.value().clone(),
|
);
|
||||||
session_token: None,
|
return ProviderConfig::Bedrock(config);
|
||||||
cross_region_inference: *settings.bedrock_cross_region_inference.value(),
|
|
||||||
};
|
|
||||||
|
|
||||||
if let ::ai::api_keys::AwsCredentialsState::Loaded { credentials, .. } =
|
|
||||||
api_key_manager.aws_credentials_state()
|
|
||||||
{
|
|
||||||
config.auth_method = crate::settings::BedrockAuthMethod::StaticKeys;
|
|
||||||
config.access_key_id = credentials.access_key().to_string();
|
|
||||||
config.secret_access_key = credentials.secret_key().to_string();
|
|
||||||
config.session_token = credentials.session_token().map(|s| s.to_string());
|
|
||||||
}
|
|
||||||
|
|
||||||
return ProviderConfig::Bedrock(config.with_external_fallbacks());
|
|
||||||
}
|
}
|
||||||
|
|
||||||
ProviderConfig::None
|
ProviderConfig::None
|
||||||
@@ -229,9 +215,8 @@ impl CrosscheckReviewer {
|
|||||||
model: config.model.unwrap_or(model_id),
|
model: config.model.unwrap_or(model_id),
|
||||||
reasoning_effort: config.reasoning_effort,
|
reasoning_effort: config.reasoning_effort,
|
||||||
max_output_tokens: Some(u64::from(REVIEWER_MAX_OUTPUT_TOKENS)),
|
max_output_tokens: Some(u64::from(REVIEWER_MAX_OUTPUT_TOKENS)),
|
||||||
auth_file: None,
|
|
||||||
},
|
},
|
||||||
);
|
).with_credential_resolver(crate::ai::chatgpt_auth::credential_resolver());
|
||||||
runtime
|
runtime
|
||||||
.complete_text(
|
.complete_text(
|
||||||
prompt::CROSSCHECK_REVIEWER_SYSTEM_PROMPT.to_string(),
|
prompt::CROSSCHECK_REVIEWER_SYSTEM_PROMPT.to_string(),
|
||||||
|
|||||||
+42
-53
@@ -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;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -3471,13 +3468,5 @@ mod tests;
|
|||||||
fn bedrock_discovery_config(
|
fn bedrock_discovery_config(
|
||||||
settings: &AISettings,
|
settings: &AISettings,
|
||||||
) -> crate::ai::provider::client::BedrockClientConfig {
|
) -> crate::ai::provider::client::BedrockClientConfig {
|
||||||
crate::ai::provider::client::BedrockClientConfig {
|
crate::ai::provider::client::BedrockClientConfig::from_settings(settings)
|
||||||
auth_method: *settings.bedrock_auth_method.value(),
|
|
||||||
profile: settings.bedrock_profile.value().clone(),
|
|
||||||
region: settings.bedrock_region.value().clone(),
|
|
||||||
access_key_id: settings.bedrock_access_key_id.value().clone(),
|
|
||||||
secret_access_key: settings.bedrock_secret_access_key.value().clone(),
|
|
||||||
session_token: None,
|
|
||||||
cross_region_inference: *settings.bedrock_cross_region_inference.value(),
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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(),
|
||||||
|
|||||||
+3
-1
@@ -85,10 +85,12 @@ pub mod outline;
|
|||||||
pub(crate) use ai::paths;
|
pub(crate) use ai::paths;
|
||||||
|
|
||||||
pub fn init(app: &mut AppContext) {
|
pub fn init(app: &mut AppContext) {
|
||||||
|
#[cfg(not(target_family = "wasm"))]
|
||||||
|
app.add_singleton_model(provider::authentication::BedrockAuthentication::new);
|
||||||
#[cfg(not(target_family = "wasm"))]
|
#[cfg(not(target_family = "wasm"))]
|
||||||
app.add_singleton_model(acp::AcpRuntimeModel::new);
|
app.add_singleton_model(acp::AcpRuntimeModel::new);
|
||||||
#[cfg(not(target_family = "wasm"))]
|
#[cfg(not(target_family = "wasm"))]
|
||||||
app.add_singleton_model(|_| chatgpt_auth::ChatGPTAuthModel::new());
|
app.add_singleton_model(chatgpt_auth::ChatGPTAuthModel::new);
|
||||||
#[cfg(not(target_family = "wasm"))]
|
#[cfg(not(target_family = "wasm"))]
|
||||||
app.add_singleton_model(claude_oauth::ClaudeOAuthModel::new);
|
app.add_singleton_model(claude_oauth::ClaudeOAuthModel::new);
|
||||||
blocklist::keyboard_navigable_buttons::init(app);
|
blocklist::keyboard_navigable_buttons::init(app);
|
||||||
|
|||||||
@@ -0,0 +1,333 @@
|
|||||||
|
use std::collections::HashMap;
|
||||||
|
use std::time::SystemTime;
|
||||||
|
|
||||||
|
use ai::api_keys::{ApiKeyManager, AwsCredentials, AwsCredentialsRefreshStrategy, AwsCredentialsState};
|
||||||
|
use futures::channel::oneshot;
|
||||||
|
use futures::future::{BoxFuture, Shared};
|
||||||
|
use futures::FutureExt;
|
||||||
|
use galaxyui::{AppContext, Entity, ModelContext, SingletonEntity};
|
||||||
|
use settings::Setting;
|
||||||
|
|
||||||
|
use super::client::{BedrockClientConfig, BedrockError};
|
||||||
|
use crate::ai::agent::RenderableAIError;
|
||||||
|
use crate::settings::{AISettings, AISettingsChangedEvent, BedrockAuthMethod};
|
||||||
|
|
||||||
|
impl BedrockClientConfig {
|
||||||
|
pub(crate) fn from_settings(settings: &AISettings) -> Self {
|
||||||
|
Self {
|
||||||
|
auth_method: *settings.bedrock_auth_method.value(),
|
||||||
|
profile: settings.bedrock_profile.value().clone(),
|
||||||
|
region: settings.bedrock_region.value().clone(),
|
||||||
|
access_key_id: settings.bedrock_access_key_id.value().clone(),
|
||||||
|
secret_access_key: settings.bedrock_secret_access_key.value().clone(),
|
||||||
|
session_token: None,
|
||||||
|
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,
|
||||||
|
state: &AwsCredentialsState,
|
||||||
|
) -> Self {
|
||||||
|
match strategy {
|
||||||
|
AwsCredentialsRefreshStrategy::LocalChain => self,
|
||||||
|
AwsCredentialsRefreshStrategy::OidcManaged { region, .. } => {
|
||||||
|
self.auth_method = BedrockAuthMethod::StaticKeys;
|
||||||
|
self.region.clone_from(region);
|
||||||
|
self.access_key_id.clear();
|
||||||
|
self.secret_access_key.clear();
|
||||||
|
self.session_token = None;
|
||||||
|
self.credentials_expiry = None;
|
||||||
|
if let AwsCredentialsState::Loaded { credentials, .. } = state {
|
||||||
|
self.access_key_id = credentials.access_key().to_owned();
|
||||||
|
self.secret_access_key = credentials.secret_key().to_owned();
|
||||||
|
self.session_token = credentials.session_token().map(str::to_owned);
|
||||||
|
self.credentials_expiry = Some(
|
||||||
|
credentials.expires_at().unwrap_or(SystemTime::UNIX_EPOCH),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
self
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn credentials_are_current(expiry: Option<SystemTime>, now: SystemTime) -> bool {
|
||||||
|
expiry.is_none_or(|expiry| expiry > now)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn provider_startup_error(error: &anyhow::Error, model: &str) -> RenderableAIError {
|
||||||
|
match error.downcast_ref::<BedrockError>() {
|
||||||
|
Some(BedrockError::CredentialsNotConfigured | BedrockError::CredentialsLoadFailed(_)) => {
|
||||||
|
RenderableAIError::AwsBedrockCredentialsExpiredOrInvalid {
|
||||||
|
model_name: model.to_owned(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Some(
|
||||||
|
BedrockError::RegionNotConfigured
|
||||||
|
| BedrockError::ApiError(_)
|
||||||
|
| BedrockError::ModelNotFound(_)
|
||||||
|
| BedrockError::AccessDenied(_)
|
||||||
|
| BedrockError::Throttling(_)
|
||||||
|
| BedrockError::ValidationError(_),
|
||||||
|
) | None => RenderableAIError::other(error.to_string(), false),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn login_command(
|
||||||
|
configured: &str,
|
||||||
|
profile: &str,
|
||||||
|
) -> String {
|
||||||
|
if configured != "aws sso login" {
|
||||||
|
return configured.to_owned();
|
||||||
|
}
|
||||||
|
if profile.is_empty() {
|
||||||
|
configured.to_owned()
|
||||||
|
} else {
|
||||||
|
format!("aws sso login --profile {}", shell_words::quote(profile))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn login_command_parts(command: &str) -> Result<Vec<String>, String> {
|
||||||
|
let parts = shell_words::split(command)
|
||||||
|
.map_err(|_| "AWS login command contains invalid quoting".to_owned())?;
|
||||||
|
if parts.is_empty() {
|
||||||
|
return Err("AWS login command is not configured".to_owned());
|
||||||
|
}
|
||||||
|
Ok(parts)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone, PartialEq)]
|
||||||
|
pub(crate) struct BedrockAuthConnection {
|
||||||
|
pub(crate) config: BedrockClientConfig,
|
||||||
|
pub(crate) enabled: bool,
|
||||||
|
pub(crate) auto_login: bool,
|
||||||
|
pub(crate) command: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl BedrockAuthConnection {
|
||||||
|
pub(crate) fn current(ctx: &AppContext) -> Self {
|
||||||
|
let settings = AISettings::as_ref(ctx);
|
||||||
|
let config = BedrockClientConfig::from_settings(settings);
|
||||||
|
Self {
|
||||||
|
command: match config.auth_method {
|
||||||
|
BedrockAuthMethod::Sso => login_command(
|
||||||
|
settings.bedrock_auth_refresh_command.value(),
|
||||||
|
&config.profile,
|
||||||
|
),
|
||||||
|
BedrockAuthMethod::Profile | BedrockAuthMethod::StaticKeys => String::new(),
|
||||||
|
},
|
||||||
|
config,
|
||||||
|
enabled: *settings.bedrock_enabled.value(),
|
||||||
|
auto_login: *settings.bedrock_auto_login.value(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn identity(&self) -> String {
|
||||||
|
format!(
|
||||||
|
"{:?}:{}:{}:{}",
|
||||||
|
self.config.auth_method, self.config.profile, self.config.region, self.command
|
||||||
|
)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Default)]
|
||||||
|
struct RecoveryAttempts {
|
||||||
|
exchanges: HashMap<String, (String, String)>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl RecoveryAttempts {
|
||||||
|
fn begin(&mut self, conversation: &str, exchange: &str, connection: &str, manual: bool) -> bool {
|
||||||
|
let attempt = (exchange.to_owned(), connection.to_owned());
|
||||||
|
if !manual && self.exchanges.get(conversation) == Some(&attempt) {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
self.exchanges.insert(conversation.to_owned(), attempt);
|
||||||
|
true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type CredentialResult = Result<AwsCredentials, String>;
|
||||||
|
type CredentialFuture = Shared<BoxFuture<'static, CredentialResult>>;
|
||||||
|
|
||||||
|
#[derive(Default)]
|
||||||
|
pub(crate) struct BedrockAuthentication {
|
||||||
|
attempts: RecoveryAttempts,
|
||||||
|
active: Vec<(u64, BedrockAuthConnection, CredentialFuture)>,
|
||||||
|
generation: u64,
|
||||||
|
revision: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl BedrockAuthentication {
|
||||||
|
pub(crate) fn new(ctx: &mut ModelContext<Self>) -> Self {
|
||||||
|
ctx.subscribe_to_model(&AISettings::handle(ctx), |me, _, event, _| {
|
||||||
|
if matches!(
|
||||||
|
event,
|
||||||
|
AISettingsChangedEvent::BedrockEnabled { .. }
|
||||||
|
| AISettingsChangedEvent::BedrockAuthMethod { .. }
|
||||||
|
| AISettingsChangedEvent::BedrockProfile { .. }
|
||||||
|
| AISettingsChangedEvent::BedrockRegion { .. }
|
||||||
|
| AISettingsChangedEvent::BedrockAuthRefreshCommand { .. }
|
||||||
|
| AISettingsChangedEvent::BedrockAccessKeyId { .. }
|
||||||
|
| AISettingsChangedEvent::BedrockSecretAccessKey { .. }
|
||||||
|
) {
|
||||||
|
me.revision += 1;
|
||||||
|
me.active.clear();
|
||||||
|
}
|
||||||
|
});
|
||||||
|
Self::default()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn revision(&self) -> u64 {
|
||||||
|
self.revision
|
||||||
|
}
|
||||||
|
pub(crate) fn recover(
|
||||||
|
&mut self,
|
||||||
|
conversation: &str,
|
||||||
|
exchange: &str,
|
||||||
|
manual: bool,
|
||||||
|
ctx: &mut ModelContext<Self>,
|
||||||
|
) -> BoxFuture<'static, Result<(), String>> {
|
||||||
|
let connection = BedrockAuthConnection::current(ctx);
|
||||||
|
if !connection.enabled || (!manual && !connection.auto_login) {
|
||||||
|
return Box::pin(async { Err("Automatic AWS login is disabled".to_owned()) });
|
||||||
|
}
|
||||||
|
if !self.attempts.begin(conversation, exchange, &connection.identity(), manual) {
|
||||||
|
return Box::pin(async {
|
||||||
|
Err("AWS authentication still failed after recovery. Use Refresh AWS Credentials to retry.".to_owned())
|
||||||
|
});
|
||||||
|
}
|
||||||
|
let strategy = ApiKeyManager::as_ref(ctx).aws_credentials_refresh_strategy();
|
||||||
|
if matches!(strategy, AwsCredentialsRefreshStrategy::OidcManaged { .. }) {
|
||||||
|
return ApiKeyManager::handle(ctx).update(ctx, |manager, ctx| {
|
||||||
|
crate::ai::aws_credentials::refresh_aws_credentials(manager, ctx)
|
||||||
|
});
|
||||||
|
}
|
||||||
|
if connection.config.auth_method == BedrockAuthMethod::StaticKeys {
|
||||||
|
return Box::pin(async {
|
||||||
|
Err("Update the configured AWS static keys to authenticate.".to_owned())
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
let (generation, future) = match self.active.iter().find(|(_, active, _)| active == &connection) {
|
||||||
|
Some((generation, _, future)) => (*generation, future.clone()),
|
||||||
|
None => {
|
||||||
|
self.generation += 1;
|
||||||
|
let config = connection.config.clone();
|
||||||
|
let command = connection.command.clone();
|
||||||
|
let future = recover_credentials(
|
||||||
|
connection.config.auth_method,
|
||||||
|
manual,
|
||||||
|
move || {
|
||||||
|
let config = config.clone();
|
||||||
|
async move {
|
||||||
|
crate::ai::aws_credentials::load_aws_credentials_from_sdk(&config)
|
||||||
|
.await
|
||||||
|
.map_err(|error| error.to_string())
|
||||||
|
}.boxed()
|
||||||
|
},
|
||||||
|
move || {
|
||||||
|
let command = command.clone();
|
||||||
|
async move {
|
||||||
|
let parts = login_command_parts(&command)?;
|
||||||
|
let status = tokio::process::Command::new(&parts[0])
|
||||||
|
.args(&parts[1..])
|
||||||
|
.status()
|
||||||
|
.await
|
||||||
|
.map_err(|error| format!("Failed to start AWS login: {error}"))?;
|
||||||
|
if !status.success() {
|
||||||
|
return Err(format!("AWS login command exited with status: {status}"));
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}.boxed()
|
||||||
|
},
|
||||||
|
).boxed().shared();
|
||||||
|
self.active.push((self.generation, connection.clone(), future.clone()));
|
||||||
|
(self.generation, future)
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let revision = self.revision;
|
||||||
|
let (tx, rx) = oneshot::channel();
|
||||||
|
let _ = ctx.spawn(future, move |me, result, ctx| {
|
||||||
|
me.active.retain(|(active, _, _)| *active != generation);
|
||||||
|
if me.revision != revision
|
||||||
|
|| BedrockAuthConnection::current(ctx) != connection
|
||||||
|
|| ApiKeyManager::as_ref(ctx).aws_credentials_refresh_strategy() != strategy
|
||||||
|
{
|
||||||
|
let _ = tx.send(Err("AWS connection changed during authentication".to_owned()));
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let outcome = result.as_ref().map(|_| ()).map_err(Clone::clone);
|
||||||
|
ApiKeyManager::handle(ctx).update(ctx, |manager, ctx| {
|
||||||
|
let state = match result {
|
||||||
|
Ok(credentials) => AwsCredentialsState::Loaded {
|
||||||
|
credentials,
|
||||||
|
loaded_at: SystemTime::now(),
|
||||||
|
},
|
||||||
|
Err(message) => AwsCredentialsState::Failed { message },
|
||||||
|
};
|
||||||
|
manager.set_aws_credentials_state(state, ctx);
|
||||||
|
});
|
||||||
|
let _ = tx.send(outcome);
|
||||||
|
});
|
||||||
|
async move {
|
||||||
|
rx.await.unwrap_or_else(|_| Err("AWS authentication was interrupted".to_owned()))
|
||||||
|
}.boxed()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn recover_credentials(
|
||||||
|
auth_method: BedrockAuthMethod,
|
||||||
|
login_first: bool,
|
||||||
|
load: impl Fn() -> BoxFuture<'static, CredentialResult>,
|
||||||
|
login: impl Fn() -> BoxFuture<'static, Result<(), String>>,
|
||||||
|
) -> CredentialResult {
|
||||||
|
match auth_method {
|
||||||
|
BedrockAuthMethod::Profile => return load().await,
|
||||||
|
BedrockAuthMethod::StaticKeys => {
|
||||||
|
return Err("Update the configured AWS static keys to authenticate.".to_owned());
|
||||||
|
}
|
||||||
|
BedrockAuthMethod::Sso => {}
|
||||||
|
}
|
||||||
|
if !login_first {
|
||||||
|
if let Ok(credentials) = load().await {
|
||||||
|
return Ok(credentials);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
login().await?;
|
||||||
|
load().await
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Entity for BedrockAuthentication {
|
||||||
|
type Event = ();
|
||||||
|
}
|
||||||
|
|
||||||
|
impl SingletonEntity for BedrockAuthentication {}
|
||||||
|
|
||||||
|
pub(crate) fn connection_revision(ctx: &AppContext) -> u64 {
|
||||||
|
ctx.try_get_singleton_model_as_ref::<BedrockAuthentication>()
|
||||||
|
.map_or(0, |authentication| authentication.revision())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
#[path = "authentication_tests.rs"]
|
||||||
|
mod tests;
|
||||||
@@ -0,0 +1,294 @@
|
|||||||
|
use std::time::{Duration, SystemTime};
|
||||||
|
|
||||||
|
use ai::api_keys::{AwsCredentials, AwsCredentialsRefreshStrategy, AwsCredentialsState};
|
||||||
|
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
fn config(auth_method: BedrockAuthMethod) -> BedrockClientConfig {
|
||||||
|
BedrockClientConfig {
|
||||||
|
auth_method,
|
||||||
|
profile: "default".to_owned(),
|
||||||
|
region: String::new(),
|
||||||
|
access_key_id: "configured-key".to_owned(),
|
||||||
|
secret_access_key: "configured-secret".to_owned(),
|
||||||
|
session_token: None,
|
||||||
|
credentials_expiry: None,
|
||||||
|
invocation_model_id: None,
|
||||||
|
cross_region_inference: true,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn provider_profile_and_region_are_authoritative_including_defaults() {
|
||||||
|
galaxyui::App::test((), |mut app| async move {
|
||||||
|
crate::test_util::settings::initialize_settings_for_tests(&mut app);
|
||||||
|
for (profile, region) in [("default", ""), ("selected-profile", "eu-west-1")] {
|
||||||
|
AISettings::handle(&app).update(&mut app, |settings, ctx| {
|
||||||
|
settings.bedrock_profile.set_value(profile.to_owned(), ctx).unwrap();
|
||||||
|
settings.bedrock_region.set_value(region.to_owned(), ctx).unwrap();
|
||||||
|
});
|
||||||
|
app.read(|ctx| {
|
||||||
|
let resolved = BedrockClientConfig::from_settings(AISettings::as_ref(ctx));
|
||||||
|
assert_eq!(resolved.profile, profile);
|
||||||
|
assert_eq!(resolved.region, region);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn custom_login_commands_are_not_rewritten() {
|
||||||
|
for command in ["", "company-login --interactive", "aws sso login --profile custom"] {
|
||||||
|
assert_eq!(login_command(command, "selected"), command);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn default_login_targets_and_quotes_the_effective_profile() {
|
||||||
|
assert_eq!(
|
||||||
|
login_command("aws sso login", "default"),
|
||||||
|
"aws sso login --profile default"
|
||||||
|
);
|
||||||
|
let command = login_command(
|
||||||
|
"aws sso login",
|
||||||
|
"team profile",
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
login_command_parts(&command).unwrap(),
|
||||||
|
["aws", "sso", "login", "--profile", "team profile"]
|
||||||
|
);
|
||||||
|
assert!(login_command_parts("").is_err());
|
||||||
|
assert!(login_command_parts("aws 'unterminated").is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn local_credential_snapshots_never_override_configured_authentication() {
|
||||||
|
let now = SystemTime::now();
|
||||||
|
for auth in [
|
||||||
|
BedrockAuthMethod::Profile,
|
||||||
|
BedrockAuthMethod::Sso,
|
||||||
|
BedrockAuthMethod::StaticKeys,
|
||||||
|
] {
|
||||||
|
for expiry in [None, Some(now - Duration::from_secs(1)), Some(now + Duration::from_secs(60))] {
|
||||||
|
let configured = config(auth);
|
||||||
|
let loaded = AwsCredentialsState::Loaded {
|
||||||
|
credentials: AwsCredentials::new(
|
||||||
|
"snapshot-key".to_owned(),
|
||||||
|
"snapshot-secret".to_owned(),
|
||||||
|
Some("snapshot-token".to_owned()),
|
||||||
|
expiry,
|
||||||
|
),
|
||||||
|
loaded_at: now,
|
||||||
|
};
|
||||||
|
assert_eq!(
|
||||||
|
configured.clone().with_managed_credentials(
|
||||||
|
&AwsCredentialsRefreshStrategy::LocalChain,
|
||||||
|
&loaded,
|
||||||
|
),
|
||||||
|
configured
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn managed_oidc_credentials_preserve_expiry_and_do_not_fall_back_to_local_keys() {
|
||||||
|
let now = SystemTime::now();
|
||||||
|
let strategy = AwsCredentialsRefreshStrategy::OidcManaged {
|
||||||
|
task_id: Some("task".to_owned()),
|
||||||
|
role_arn: "test-role".to_owned(),
|
||||||
|
region: "us-east-1".to_owned(),
|
||||||
|
};
|
||||||
|
let loaded = AwsCredentialsState::Loaded {
|
||||||
|
credentials: AwsCredentials::new(
|
||||||
|
"managed-key".to_owned(),
|
||||||
|
"managed-secret".to_owned(),
|
||||||
|
Some("managed-token".to_owned()),
|
||||||
|
Some(now + Duration::from_secs(60)),
|
||||||
|
),
|
||||||
|
loaded_at: now,
|
||||||
|
};
|
||||||
|
let resolved = config(BedrockAuthMethod::Profile).with_managed_credentials(&strategy, &loaded);
|
||||||
|
assert_eq!(resolved.auth_method, BedrockAuthMethod::StaticKeys);
|
||||||
|
assert_eq!(resolved.region, "us-east-1");
|
||||||
|
assert_eq!(resolved.access_key_id, "managed-key");
|
||||||
|
assert_eq!(resolved.credentials_expiry, Some(now + Duration::from_secs(60)));
|
||||||
|
for state in [
|
||||||
|
AwsCredentialsState::Missing,
|
||||||
|
AwsCredentialsState::Disabled,
|
||||||
|
AwsCredentialsState::Refreshing,
|
||||||
|
AwsCredentialsState::Failed { message: "failed".to_owned() },
|
||||||
|
] {
|
||||||
|
let resolved = config(BedrockAuthMethod::Profile).with_managed_credentials(&strategy, &state);
|
||||||
|
assert!(resolved.access_key_id.is_empty());
|
||||||
|
assert!(resolved.secret_access_key.is_empty());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn automatic_recovery_is_bounded_per_user_exchange_but_manual_retry_is_available() {
|
||||||
|
let mut attempts = RecoveryAttempts::default();
|
||||||
|
assert!(attempts.begin("conversation", "query-1", "connection", false));
|
||||||
|
assert!(!attempts.begin("conversation", "query-1", "connection", false));
|
||||||
|
assert!(attempts.begin("conversation", "query-1", "connection", true));
|
||||||
|
assert!(attempts.begin("conversation", "query-2", "connection", false));
|
||||||
|
assert!(attempts.begin("conversation", "query-2", "new-connection", false));
|
||||||
|
assert!(attempts.begin("another-conversation", "query-2", "connection", false));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn credential_expiry_is_checked_at_the_boundary() {
|
||||||
|
let now = SystemTime::now();
|
||||||
|
assert!(credentials_are_current(None, now));
|
||||||
|
assert!(credentials_are_current(Some(now + Duration::from_secs(1)), now));
|
||||||
|
assert!(!credentials_are_current(Some(now), 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 {
|
||||||
|
AwsCredentials::new("test-key".to_owned(), "test-secret".to_owned(), None, None)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn external_refresh_recovers_without_running_login() {
|
||||||
|
let credentials = recover_credentials(
|
||||||
|
BedrockAuthMethod::Sso,
|
||||||
|
false,
|
||||||
|
|| futures::future::ready(Ok(test_credentials())).boxed(),
|
||||||
|
|| futures::future::ready(Err("Login must not run".to_owned())).boxed(),
|
||||||
|
).await.unwrap();
|
||||||
|
assert_eq!(credentials, test_credentials());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn expired_session_runs_login_then_awaits_a_fresh_reload() {
|
||||||
|
use std::sync::{Arc, atomic::{AtomicUsize, Ordering}};
|
||||||
|
|
||||||
|
let calls = Arc::new(AtomicUsize::new(0));
|
||||||
|
let load_calls = calls.clone();
|
||||||
|
let login_calls = calls.clone();
|
||||||
|
let result = recover_credentials(
|
||||||
|
BedrockAuthMethod::Sso,
|
||||||
|
false,
|
||||||
|
move || {
|
||||||
|
let index = load_calls.fetch_add(1, Ordering::SeqCst);
|
||||||
|
let result = match index {
|
||||||
|
0 => Err("expired".to_owned()),
|
||||||
|
2 => Ok(test_credentials()),
|
||||||
|
index => panic!("Unexpected reload order: {index}"),
|
||||||
|
};
|
||||||
|
futures::future::ready(result).boxed()
|
||||||
|
},
|
||||||
|
move || {
|
||||||
|
assert_eq!(login_calls.fetch_add(1, Ordering::SeqCst), 1);
|
||||||
|
futures::future::ready(Ok(())).boxed()
|
||||||
|
},
|
||||||
|
).await;
|
||||||
|
assert_eq!(result, Ok(test_credentials()));
|
||||||
|
assert_eq!(calls.load(Ordering::SeqCst), 3);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn manual_login_failure_does_not_reload_or_resume() {
|
||||||
|
let result = recover_credentials(
|
||||||
|
BedrockAuthMethod::Sso,
|
||||||
|
true,
|
||||||
|
|| panic!("Credentials must not reload after failed login"),
|
||||||
|
|| futures::future::ready(Err("login failed".to_owned())).boxed(),
|
||||||
|
).await;
|
||||||
|
assert_eq!(result, Err("login failed".to_owned()));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn login_success_with_failed_reload_remains_an_error() {
|
||||||
|
let result = recover_credentials(
|
||||||
|
BedrockAuthMethod::Sso,
|
||||||
|
true,
|
||||||
|
|| futures::future::ready(Err("reload failed".to_owned())).boxed(),
|
||||||
|
|| futures::future::ready(Ok(())).boxed(),
|
||||||
|
).await;
|
||||||
|
assert_eq!(result, Err("reload failed".to_owned()));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn shared_recovery_executes_login_once_for_overlapping_waiters() {
|
||||||
|
use std::sync::{Arc, atomic::{AtomicUsize, Ordering}};
|
||||||
|
|
||||||
|
let calls = Arc::new(AtomicUsize::new(0));
|
||||||
|
let login_calls = calls.clone();
|
||||||
|
let recovery = recover_credentials(
|
||||||
|
BedrockAuthMethod::Sso,
|
||||||
|
true,
|
||||||
|
|| futures::future::ready(Ok(test_credentials())).boxed(),
|
||||||
|
move || {
|
||||||
|
login_calls.fetch_add(1, Ordering::SeqCst);
|
||||||
|
futures::future::ready(Ok(())).boxed()
|
||||||
|
},
|
||||||
|
).boxed().shared();
|
||||||
|
let (first, second) = futures::join!(recovery.clone(), recovery);
|
||||||
|
assert_eq!(first, Ok(test_credentials()));
|
||||||
|
assert_eq!(second, Ok(test_credentials()));
|
||||||
|
assert_eq!(calls.load(Ordering::SeqCst), 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn credential_preparation_failures_surface_authentication_recovery() {
|
||||||
|
use super::super::client::BedrockError;
|
||||||
|
|
||||||
|
for error in [
|
||||||
|
BedrockError::CredentialsNotConfigured,
|
||||||
|
BedrockError::CredentialsLoadFailed("SSO session invalid".to_owned()),
|
||||||
|
] {
|
||||||
|
let error = provider_startup_error(&anyhow::Error::new(error), "selected-model");
|
||||||
|
assert!(error.is_aws_bedrock_credentials_error());
|
||||||
|
}
|
||||||
|
for error in [
|
||||||
|
BedrockError::RegionNotConfigured,
|
||||||
|
BedrockError::ApiError("transport".to_owned()),
|
||||||
|
BedrockError::ModelNotFound("model".to_owned()),
|
||||||
|
BedrockError::AccessDenied("policy".to_owned()),
|
||||||
|
BedrockError::Throttling("busy".to_owned()),
|
||||||
|
BedrockError::ValidationError("request".to_owned()),
|
||||||
|
] {
|
||||||
|
let error = provider_startup_error(&anyhow::Error::new(error), "selected-model");
|
||||||
|
assert!(!error.is_aws_bedrock_credentials_error());
|
||||||
|
}
|
||||||
|
assert!(!provider_startup_error(&anyhow::anyhow!("unrelated"), "selected-model")
|
||||||
|
.is_aws_bedrock_credentials_error());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn profile_authentication_never_starts_a_login_flow() {
|
||||||
|
for manual in [false, true] {
|
||||||
|
let result = recover_credentials(
|
||||||
|
BedrockAuthMethod::Profile,
|
||||||
|
manual,
|
||||||
|
|| futures::future::ready(Ok(test_credentials())).boxed(),
|
||||||
|
|| panic!("Profile authentication must not execute login commands"),
|
||||||
|
).await;
|
||||||
|
assert_eq!(result, Ok(test_credentials()));
|
||||||
|
let result = recover_credentials(
|
||||||
|
BedrockAuthMethod::Profile,
|
||||||
|
manual,
|
||||||
|
|| futures::future::ready(Err("profile needs fresh credentials".to_owned())).boxed(),
|
||||||
|
|| panic!("A failed profile reload must not execute login commands"),
|
||||||
|
).await;
|
||||||
|
assert_eq!(result, Err("profile needs fresh credentials".to_owned()));
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -1,4 +1,5 @@
|
|||||||
use std::sync::{Arc, Mutex};
|
use std::sync::{Arc, Mutex};
|
||||||
|
use std::time::SystemTime;
|
||||||
|
|
||||||
use anyhow::Result;
|
use anyhow::Result;
|
||||||
use aws_config::BehaviorVersion;
|
use aws_config::BehaviorVersion;
|
||||||
@@ -27,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)]
|
||||||
@@ -37,40 +39,18 @@ pub struct BedrockClientConfig {
|
|||||||
pub access_key_id: String,
|
pub access_key_id: String,
|
||||||
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 invocation_model_id: Option<String>,
|
||||||
pub cross_region_inference: bool,
|
pub cross_region_inference: bool,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl BedrockClientConfig {
|
|
||||||
/// Applies external config (from Claude Code / OpenCode) as fallback values
|
|
||||||
/// when Galaxy's own settings are at their defaults.
|
|
||||||
pub fn with_external_fallbacks(mut self) -> Self {
|
|
||||||
let external = ExternalBedrockConfig::load();
|
|
||||||
if external.is_empty() {
|
|
||||||
return self;
|
|
||||||
}
|
|
||||||
|
|
||||||
if self.profile == "default" {
|
|
||||||
if let Some(profile) = external.profile {
|
|
||||||
log::info!("[bedrock] Using profile from external config: {profile}");
|
|
||||||
self.profile = profile;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
if self.region.is_empty() {
|
|
||||||
if let Some(region) = external.region {
|
|
||||||
log::info!("[bedrock] Using region from external config: {region}");
|
|
||||||
self.region = region;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
self
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(Debug, thiserror::Error)]
|
#[derive(Debug, thiserror::Error)]
|
||||||
pub enum BedrockError {
|
pub enum BedrockError {
|
||||||
#[error("Bedrock credentials not configured")]
|
#[error("Bedrock credentials not configured")]
|
||||||
CredentialsNotConfigured,
|
CredentialsNotConfigured,
|
||||||
|
#[error("Failed to load AWS credentials: {0}")]
|
||||||
|
CredentialsLoadFailed(String),
|
||||||
#[error("Bedrock region not configured and could not be auto-detected")]
|
#[error("Bedrock region not configured and could not be auto-detected")]
|
||||||
RegionNotConfigured,
|
RegionNotConfigured,
|
||||||
#[error("Bedrock API error: {0}")]
|
#[error("Bedrock API error: {0}")]
|
||||||
@@ -115,12 +95,20 @@ impl BedrockClient {
|
|||||||
if config.access_key_id.is_empty() || config.secret_access_key.is_empty() {
|
if config.access_key_id.is_empty() || config.secret_access_key.is_empty() {
|
||||||
return Err(BedrockError::CredentialsNotConfigured);
|
return Err(BedrockError::CredentialsNotConfigured);
|
||||||
}
|
}
|
||||||
|
if !super::authentication::credentials_are_current(
|
||||||
|
config.credentials_expiry,
|
||||||
|
SystemTime::now(),
|
||||||
|
) {
|
||||||
|
return Err(BedrockError::CredentialsLoadFailed(
|
||||||
|
"AWS credentials have expired".to_owned(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
let creds = aws_credential_types::Credentials::new(
|
let creds = aws_credential_types::Credentials::new(
|
||||||
&config.access_key_id,
|
&config.access_key_id,
|
||||||
&config.secret_access_key,
|
&config.secret_access_key,
|
||||||
config.session_token,
|
config.session_token,
|
||||||
None,
|
config.credentials_expiry,
|
||||||
"warp-bedrock-static",
|
"warp-bedrock-static",
|
||||||
);
|
);
|
||||||
|
|
||||||
@@ -148,11 +136,15 @@ impl BedrockClient {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
log::warn!("[bedrock] Failed to resolve AWS credentials from provider: {e:?}");
|
let message = crate::ai::aws_credentials::user_facing_aws_credentials_error_message(
|
||||||
|
&e,
|
||||||
|
&config.profile,
|
||||||
|
);
|
||||||
|
return Err(BedrockError::CredentialsLoadFailed(message));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
log::warn!("[bedrock] No credentials provider found in resolved AWS config");
|
return Err(BedrockError::CredentialsNotConfigured);
|
||||||
}
|
}
|
||||||
|
|
||||||
let region = aws_config
|
let region = aws_config
|
||||||
@@ -165,6 +157,7 @@ impl BedrockClient {
|
|||||||
Ok(Self {
|
Ok(Self {
|
||||||
runtime_client,
|
runtime_client,
|
||||||
region,
|
region,
|
||||||
|
invocation_model_id: config.invocation_model_id,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -195,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)");
|
||||||
|
|||||||
@@ -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 {
|
||||||
|
|||||||
@@ -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;
|
||||||
|
|
||||||
|
|||||||
@@ -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
|
||||||
@@ -161,7 +304,7 @@ async fn load_aws_config(
|
|||||||
&config.access_key_id,
|
&config.access_key_id,
|
||||||
&config.secret_access_key,
|
&config.secret_access_key,
|
||||||
config.session_token.clone(),
|
config.session_token.clone(),
|
||||||
None,
|
config.credentials_expiry,
|
||||||
"galaxy-bedrock-discovery",
|
"galaxy-bedrock-discovery",
|
||||||
);
|
);
|
||||||
let mut loader =
|
let mut loader =
|
||||||
|
|||||||
@@ -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());
|
||||||
|
|||||||
@@ -268,6 +268,8 @@ fn get_test_config() -> Option<BedrockClientConfig> {
|
|||||||
access_key_id: String::new(),
|
access_key_id: String::new(),
|
||||||
secret_access_key: String::new(),
|
secret_access_key: String::new(),
|
||||||
session_token: None,
|
session_token: None,
|
||||||
|
credentials_expiry: None,
|
||||||
|
invocation_model_id: None,
|
||||||
cross_region_inference: false,
|
cross_region_inference: false,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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()
|
||||||
|
|||||||
@@ -23,6 +23,8 @@ fn get_test_config() -> Option<BedrockClientConfig> {
|
|||||||
access_key_id: String::new(),
|
access_key_id: String::new(),
|
||||||
secret_access_key: String::new(),
|
secret_access_key: String::new(),
|
||||||
session_token: None,
|
session_token: None,
|
||||||
|
credentials_expiry: None,
|
||||||
|
invocation_model_id: None,
|
||||||
cross_region_inference: false,
|
cross_region_inference: false,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -16,6 +16,8 @@ pub enum ProviderConfig {
|
|||||||
// Provider-neutral module boundary; Bedrock-specific transport helpers live here without a
|
// Provider-neutral module boundary; Bedrock-specific transport helpers live here without a
|
||||||
// provider-named directory so the provider layer can be reorganized independently of callers.
|
// provider-named directory so the provider layer can be reorganized independently of callers.
|
||||||
#[cfg(not(target_family = "wasm"))]
|
#[cfg(not(target_family = "wasm"))]
|
||||||
|
pub(crate) mod authentication;
|
||||||
|
#[cfg(not(target_family = "wasm"))]
|
||||||
pub mod client;
|
pub mod client;
|
||||||
#[cfg(not(target_family = "wasm"))]
|
#[cfg(not(target_family = "wasm"))]
|
||||||
pub mod convert;
|
pub mod convert;
|
||||||
|
|||||||
@@ -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();
|
||||||
|
|||||||
@@ -314,9 +314,8 @@ pub(crate) async fn provider_runtime_for_request(
|
|||||||
model,
|
model,
|
||||||
reasoning_effort: config.reasoning_effort,
|
reasoning_effort: config.reasoning_effort,
|
||||||
max_output_tokens: config.max_output_tokens.map(u64::from),
|
max_output_tokens: config.max_output_tokens.map(u64::from),
|
||||||
auth_file: None,
|
|
||||||
},
|
},
|
||||||
)),
|
).with_credential_resolver(crate::ai::chatgpt_auth::credential_resolver())),
|
||||||
OpenAIProviderKind::Anthropic => {
|
OpenAIProviderKind::Anthropic => {
|
||||||
let credential = config.api_key.unwrap_or_default();
|
let credential = config.api_key.unwrap_or_default();
|
||||||
let stored = crate::ai::claude_subscription::is_stored_credential(&credential);
|
let stored = crate::ai::claude_subscription::is_stored_credential(&credential);
|
||||||
@@ -346,14 +345,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,
|
||||||
|
|||||||
@@ -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
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -43,7 +43,7 @@ use strum::IntoEnumIterator;
|
|||||||
use super::execution_profile_view::{ExecutionProfileView, ExecutionProfileViewEvent};
|
use super::execution_profile_view::{ExecutionProfileView, ExecutionProfileViewEvent};
|
||||||
use super::provider_setup_view::{
|
use super::provider_setup_view::{
|
||||||
model_catalog_table, AcpProviderDraft, BedrockProviderDraft, ProviderSetupProviderType,
|
model_catalog_table, AcpProviderDraft, BedrockProviderDraft, ProviderSetupProviderType,
|
||||||
ProviderSetupView, ProviderSetupViewEvent,
|
ProviderSetupView, ProviderSetupViewAction, ProviderSetupViewEvent,
|
||||||
};
|
};
|
||||||
use super::settings_page::{
|
use super::settings_page::{
|
||||||
build_sub_header, build_toggle_element, render_body_item_label,
|
build_sub_header, build_toggle_element, render_body_item_label,
|
||||||
@@ -2720,6 +2720,7 @@ pub enum AISettingsPageAction {
|
|||||||
},
|
},
|
||||||
RemoveOpenAIProvider(usize),
|
RemoveOpenAIProvider(usize),
|
||||||
EditBedrockProvider,
|
EditBedrockProvider,
|
||||||
|
RefreshBedrockModels,
|
||||||
RemoveBedrockProvider,
|
RemoveBedrockProvider,
|
||||||
RemoveAcpProvider(String),
|
RemoveAcpProvider(String),
|
||||||
ToggleFileBasedMcp,
|
ToggleFileBasedMcp,
|
||||||
@@ -3457,6 +3458,12 @@ impl TypedActionView for AISettingsPageView {
|
|||||||
AISettingsPageAction::EditBedrockProvider => {
|
AISettingsPageAction::EditBedrockProvider => {
|
||||||
self.begin_inline_bedrock_setup(ctx);
|
self.begin_inline_bedrock_setup(ctx);
|
||||||
}
|
}
|
||||||
|
AISettingsPageAction::RefreshBedrockModels => {
|
||||||
|
self.begin_inline_bedrock_setup(ctx);
|
||||||
|
self.provider_setup_body.update(ctx, |body, ctx| {
|
||||||
|
body.handle_action(&ProviderSetupViewAction::Next, ctx);
|
||||||
|
});
|
||||||
|
}
|
||||||
AISettingsPageAction::RemoveBedrockProvider => {
|
AISettingsPageAction::RemoveBedrockProvider => {
|
||||||
AISettings::handle(ctx).update(ctx, |settings, ctx| {
|
AISettings::handle(ctx).update(ctx, |settings, ctx| {
|
||||||
report_if_error!(settings.bedrock_enabled.set_value(false, ctx));
|
report_if_error!(settings.bedrock_enabled.set_value(false, ctx));
|
||||||
@@ -6641,6 +6648,7 @@ struct ProviderSettingsWidget {
|
|||||||
bedrock_add_button: ViewHandle<ActionButton>,
|
bedrock_add_button: ViewHandle<ActionButton>,
|
||||||
acp_add_button: ViewHandle<ActionButton>,
|
acp_add_button: ViewHandle<ActionButton>,
|
||||||
bedrock_edit_button: ViewHandle<ActionButton>,
|
bedrock_edit_button: ViewHandle<ActionButton>,
|
||||||
|
bedrock_refresh_button: ViewHandle<ActionButton>,
|
||||||
bedrock_remove_button: ViewHandle<ActionButton>,
|
bedrock_remove_button: ViewHandle<ActionButton>,
|
||||||
provider_cards: Vec<OpenAIProviderCardState>,
|
provider_cards: Vec<OpenAIProviderCardState>,
|
||||||
acp_provider_cards: Vec<AcpProviderCardState>,
|
acp_provider_cards: Vec<AcpProviderCardState>,
|
||||||
@@ -6749,6 +6757,11 @@ impl ProviderSettingsWidget {
|
|||||||
ctx.dispatch_typed_action(AISettingsPageAction::EditBedrockProvider);
|
ctx.dispatch_typed_action(AISettingsPageAction::EditBedrockProvider);
|
||||||
})
|
})
|
||||||
});
|
});
|
||||||
|
let bedrock_refresh_button = ctx.add_typed_action_view(|_| {
|
||||||
|
ActionButton::new("Refresh models", SecondaryTheme).on_click(|ctx| {
|
||||||
|
ctx.dispatch_typed_action(AISettingsPageAction::RefreshBedrockModels);
|
||||||
|
})
|
||||||
|
});
|
||||||
let bedrock_remove_button = ctx.add_typed_action_view(|_| {
|
let bedrock_remove_button = ctx.add_typed_action_view(|_| {
|
||||||
ActionButton::new("Delete", DangerSecondaryTheme).on_click(|ctx| {
|
ActionButton::new("Delete", DangerSecondaryTheme).on_click(|ctx| {
|
||||||
ctx.dispatch_typed_action(AISettingsPageAction::RemoveBedrockProvider);
|
ctx.dispatch_typed_action(AISettingsPageAction::RemoveBedrockProvider);
|
||||||
@@ -6766,6 +6779,7 @@ impl ProviderSettingsWidget {
|
|||||||
bedrock_add_button,
|
bedrock_add_button,
|
||||||
acp_add_button,
|
acp_add_button,
|
||||||
bedrock_edit_button,
|
bedrock_edit_button,
|
||||||
|
bedrock_refresh_button,
|
||||||
bedrock_remove_button,
|
bedrock_remove_button,
|
||||||
provider_cards,
|
provider_cards,
|
||||||
acp_provider_cards,
|
acp_provider_cards,
|
||||||
@@ -7275,6 +7289,7 @@ impl ProviderSettingsWidget {
|
|||||||
.with_child(
|
.with_child(
|
||||||
Flex::row()
|
Flex::row()
|
||||||
.with_spacing(8.)
|
.with_spacing(8.)
|
||||||
|
.with_child(ChildView::new(&self.bedrock_refresh_button).finish())
|
||||||
.with_child(ChildView::new(&self.bedrock_edit_button).finish())
|
.with_child(ChildView::new(&self.bedrock_edit_button).finish())
|
||||||
.with_child(ChildView::new(&self.bedrock_remove_button).finish())
|
.with_child(ChildView::new(&self.bedrock_remove_button).finish())
|
||||||
.finish(),
|
.finish(),
|
||||||
|
|||||||
@@ -220,6 +220,7 @@ pub enum ProviderSetupViewAction {
|
|||||||
ToggleModel(usize),
|
ToggleModel(usize),
|
||||||
CycleModelCapability(usize, CapabilityKey),
|
CycleModelCapability(usize, CapabilityKey),
|
||||||
ConnectChatGPT,
|
ConnectChatGPT,
|
||||||
|
CancelChatGPTSignIn,
|
||||||
ConnectClaude,
|
ConnectClaude,
|
||||||
SignOutClaude,
|
SignOutClaude,
|
||||||
CancelClaudeSignIn,
|
CancelClaudeSignIn,
|
||||||
@@ -687,10 +688,8 @@ impl ProviderSetupView {
|
|||||||
self.draft_bedrock = draft;
|
self.draft_bedrock = draft;
|
||||||
self.draft_models.clear();
|
self.draft_models.clear();
|
||||||
self.discovery_state = DiscoveryState::Idle;
|
self.discovery_state = DiscoveryState::Idle;
|
||||||
|
self.bedrock_region_request += 1;
|
||||||
self.sync_editors(ctx);
|
self.sync_editors(ctx);
|
||||||
if self.draft_bedrock.region.is_empty() {
|
|
||||||
self.load_bedrock_profile_region(ctx);
|
|
||||||
}
|
|
||||||
self.sync_provider_type_buttons(ctx);
|
self.sync_provider_type_buttons(ctx);
|
||||||
self.sync_bedrock_auth_buttons(ctx);
|
self.sync_bedrock_auth_buttons(ctx);
|
||||||
self.sync_model_switches(ctx);
|
self.sync_model_switches(ctx);
|
||||||
@@ -1130,8 +1129,23 @@ 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.bedrock_region_request += 1;
|
||||||
self.step = ProviderSetupStep::Discover;
|
self.step = ProviderSetupStep::Discover;
|
||||||
self.discovery_state = DiscoveryState::Loading;
|
self.discovery_state = DiscoveryState::Loading;
|
||||||
self.acp_setup_status.clear();
|
self.acp_setup_status.clear();
|
||||||
@@ -1145,33 +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,
|
|
||||||
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,
|
||||||
@@ -1189,6 +1215,7 @@ impl ProviderSetupView {
|
|||||||
me.update_next_button(ctx);
|
me.update_next_button(ctx);
|
||||||
ctx.notify();
|
ctx.notify();
|
||||||
}
|
}
|
||||||
|
}
|
||||||
},
|
},
|
||||||
);
|
);
|
||||||
return;
|
return;
|
||||||
@@ -1564,10 +1591,10 @@ impl ProviderSetupView {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
if matches!(
|
let pending = matches!(state,
|
||||||
state,
|
ChatGPTAuthState::AwaitingBrowser | ChatGPTAuthState::ExchangingToken
|
||||||
ChatGPTAuthState::NotConnected | ChatGPTAuthState::Failed(_)
|
);
|
||||||
) {
|
{
|
||||||
children.push(
|
children.push(
|
||||||
appearance
|
appearance
|
||||||
.ui_builder()
|
.ui_builder()
|
||||||
@@ -1575,10 +1602,21 @@ impl ProviderSetupView {
|
|||||||
ButtonVariant::Secondary,
|
ButtonVariant::Secondary,
|
||||||
self.chatgpt_connect_mouse_state.clone(),
|
self.chatgpt_connect_mouse_state.clone(),
|
||||||
)
|
)
|
||||||
.with_text_label("Connect ChatGPT".to_owned())
|
.with_text_label(if pending {
|
||||||
|
"Cancel sign-in"
|
||||||
|
} else if matches!(state, ChatGPTAuthState::Connected) {
|
||||||
|
"Reconnect ChatGPT"
|
||||||
|
} else {
|
||||||
|
"Connect ChatGPT"
|
||||||
|
}.to_owned())
|
||||||
.build()
|
.build()
|
||||||
.on_click(|ctx, _, _| {
|
.on_click(move |ctx, _, _| {
|
||||||
ctx.dispatch_typed_action(ProviderSetupViewAction::ConnectChatGPT);
|
let action = if pending {
|
||||||
|
ProviderSetupViewAction::CancelChatGPTSignIn
|
||||||
|
} else {
|
||||||
|
ProviderSetupViewAction::ConnectChatGPT
|
||||||
|
};
|
||||||
|
ctx.dispatch_typed_action(action);
|
||||||
})
|
})
|
||||||
.finish(),
|
.finish(),
|
||||||
);
|
);
|
||||||
@@ -1835,7 +1873,7 @@ impl ProviderSetupView {
|
|||||||
.finish()
|
.finish()
|
||||||
}
|
}
|
||||||
|
|
||||||
fn render_discovery(&self, appearance: &Appearance) -> Box<dyn galaxyui::Element> {
|
fn render_discovery(&self, appearance: &Appearance, app: &AppContext) -> Box<dyn galaxyui::Element> {
|
||||||
let (message, color) = match &self.discovery_state {
|
let (message, color) = match &self.discovery_state {
|
||||||
DiscoveryState::Idle | DiscoveryState::Loading => (
|
DiscoveryState::Idle | DiscoveryState::Loading => (
|
||||||
if self.provider_type == ProviderSetupProviderType::Acp
|
if self.provider_type == ProviderSetupProviderType::Acp
|
||||||
@@ -1849,7 +1887,7 @@ impl ProviderSetupView {
|
|||||||
),
|
),
|
||||||
DiscoveryState::Failed(error) => (error.clone(), appearance.theme().ui_error_color()),
|
DiscoveryState::Failed(error) => (error.clone(), appearance.theme().ui_error_color()),
|
||||||
};
|
};
|
||||||
Flex::column()
|
let mut column = Flex::column()
|
||||||
.with_spacing(12.)
|
.with_spacing(12.)
|
||||||
.with_child(
|
.with_child(
|
||||||
Text::new(
|
Text::new(
|
||||||
@@ -1887,8 +1925,13 @@ impl ProviderSetupView {
|
|||||||
.with_color(color)
|
.with_color(color)
|
||||||
.soft_wrap(true)
|
.soft_wrap(true)
|
||||||
.finish()
|
.finish()
|
||||||
})
|
});
|
||||||
.finish()
|
if self.provider_type == ProviderSetupProviderType::ChatGPTSubscription
|
||||||
|
&& matches!(self.discovery_state, DiscoveryState::Failed(_))
|
||||||
|
{
|
||||||
|
column = column.with_child(self.render_chatgpt_auth(appearance, app));
|
||||||
|
}
|
||||||
|
column.finish()
|
||||||
}
|
}
|
||||||
|
|
||||||
fn render_model_table(
|
fn render_model_table(
|
||||||
@@ -2442,7 +2485,7 @@ impl View for ProviderSetupView {
|
|||||||
let content = match self.step {
|
let content = match self.step {
|
||||||
ProviderSetupStep::ProviderType => self.render_provider_type(appearance),
|
ProviderSetupStep::ProviderType => self.render_provider_type(appearance),
|
||||||
ProviderSetupStep::Configure => self.render_configure(appearance, app),
|
ProviderSetupStep::Configure => self.render_configure(appearance, app),
|
||||||
ProviderSetupStep::Discover => self.render_discovery(appearance),
|
ProviderSetupStep::Discover => self.render_discovery(appearance, app),
|
||||||
ProviderSetupStep::Models => self.render_models(appearance),
|
ProviderSetupStep::Models => self.render_models(appearance),
|
||||||
};
|
};
|
||||||
Flex::column()
|
Flex::column()
|
||||||
@@ -2462,6 +2505,10 @@ impl TypedActionView for ProviderSetupView {
|
|||||||
match action {
|
match action {
|
||||||
ProviderSetupViewAction::SelectProvider(kind) => {
|
ProviderSetupViewAction::SelectProvider(kind) => {
|
||||||
if self.provider_type != *kind {
|
if self.provider_type != *kind {
|
||||||
|
#[cfg(not(target_family = "wasm"))]
|
||||||
|
if self.provider_type == ProviderSetupProviderType::ChatGPTSubscription {
|
||||||
|
ChatGPTAuthModel::handle(ctx).update(ctx, |model, ctx| model.cancel_sign_in(ctx));
|
||||||
|
}
|
||||||
self.cancel_claude_signin();
|
self.cancel_claude_signin();
|
||||||
self.draft_api_key = None;
|
self.draft_api_key = None;
|
||||||
self.claude_credentials_available = false;
|
self.claude_credentials_available = false;
|
||||||
@@ -2562,6 +2609,10 @@ impl TypedActionView for ProviderSetupView {
|
|||||||
ProviderSetupViewAction::Cancel => {
|
ProviderSetupViewAction::Cancel => {
|
||||||
self.cancel_acp_setup();
|
self.cancel_acp_setup();
|
||||||
self.cancel_claude_signin();
|
self.cancel_claude_signin();
|
||||||
|
#[cfg(not(target_family = "wasm"))]
|
||||||
|
if self.provider_type == ProviderSetupProviderType::ChatGPTSubscription {
|
||||||
|
ChatGPTAuthModel::handle(ctx).update(ctx, |model, ctx| model.cancel_sign_in(ctx));
|
||||||
|
}
|
||||||
self.discovery_state = DiscoveryState::Idle;
|
self.discovery_state = DiscoveryState::Idle;
|
||||||
ctx.emit(ProviderSetupViewEvent::Close);
|
ctx.emit(ProviderSetupViewEvent::Close);
|
||||||
}
|
}
|
||||||
@@ -2607,6 +2658,10 @@ impl TypedActionView for ProviderSetupView {
|
|||||||
#[cfg(not(target_family = "wasm"))]
|
#[cfg(not(target_family = "wasm"))]
|
||||||
ChatGPTAuthModel::handle(ctx).update(ctx, |model, ctx| model.connect(ctx));
|
ChatGPTAuthModel::handle(ctx).update(ctx, |model, ctx| model.connect(ctx));
|
||||||
}
|
}
|
||||||
|
ProviderSetupViewAction::CancelChatGPTSignIn => {
|
||||||
|
#[cfg(not(target_family = "wasm"))]
|
||||||
|
ChatGPTAuthModel::handle(ctx).update(ctx, |model, ctx| model.cancel_sign_in(ctx));
|
||||||
|
}
|
||||||
ProviderSetupViewAction::ConnectClaude => {
|
ProviderSetupViewAction::ConnectClaude => {
|
||||||
#[cfg(not(target_family = "wasm"))]
|
#[cfg(not(target_family = "wasm"))]
|
||||||
self.connect_claude(ctx);
|
self.connect_claude(ctx);
|
||||||
|
|||||||
+47
-103
@@ -29,6 +29,8 @@ use warp_util::remote_path::RemotePath;
|
|||||||
use warp_util::standardized_path::StandardizedPath;
|
use warp_util::standardized_path::StandardizedPath;
|
||||||
|
|
||||||
use crate::ai::block_context::BlockContext;
|
use crate::ai::block_context::BlockContext;
|
||||||
|
#[cfg(not(target_family = "wasm"))]
|
||||||
|
use crate::ai::provider::authentication::{BedrockAuthConnection, BedrockAuthentication};
|
||||||
use crate::global_resource_handles::GlobalResourceHandlesProvider;
|
use crate::global_resource_handles::GlobalResourceHandlesProvider;
|
||||||
pub(crate) mod docker_sandbox;
|
pub(crate) mod docker_sandbox;
|
||||||
mod link_detection;
|
mod link_detection;
|
||||||
@@ -10910,89 +10912,55 @@ impl TerminalView {
|
|||||||
self.remove_aws_bedrock_login_banner(ctx);
|
self.remove_aws_bedrock_login_banner(ctx);
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Runs the AWS login command configured in settings to refresh Bedrock credentials.
|
/// Refreshes the configured AWS connection and resumes only its originating conversation.
|
||||||
/// Doing this in PTY vs just a subprocess allows the user to see any output/errors
|
|
||||||
/// from the command directly in the terminal. Also, `aws login` commands may require
|
|
||||||
/// user interaction (e.g. "do you want to override X profile? y/n" is common)
|
|
||||||
fn run_aws_login_command(&mut self, ctx: &mut ViewContext<Self>) {
|
fn run_aws_login_command(&mut self, ctx: &mut ViewContext<Self>) {
|
||||||
let settings_command = AISettings::as_ref(ctx)
|
let conversation_id = if FeatureFlag::AgentView.is_enabled() {
|
||||||
.bedrock_auth_refresh_command
|
self.agent_view_controller
|
||||||
.value()
|
.as_ref(ctx)
|
||||||
.clone();
|
.agent_view_state()
|
||||||
|
.active_conversation_id()
|
||||||
// Use the configured command, but if it's the bare default ("aws sso login")
|
|
||||||
// prefer the external config's auth_refresh_command which includes --profile.
|
|
||||||
let login_command = if settings_command == "aws sso login" {
|
|
||||||
use crate::ai::provider::external_config::ExternalBedrockConfig;
|
|
||||||
let external = ExternalBedrockConfig::load();
|
|
||||||
external.auth_refresh_command.unwrap_or(settings_command)
|
|
||||||
} else {
|
} else {
|
||||||
settings_command
|
BlocklistAIHistoryModel::as_ref(ctx).last_conversation_id(self.id())
|
||||||
};
|
};
|
||||||
|
self.recover_aws_authentication(conversation_id, true, ctx);
|
||||||
|
}
|
||||||
|
|
||||||
if login_command.is_empty() {
|
fn recover_aws_authentication(
|
||||||
log::warn!("AWS login command is not configured");
|
&mut self,
|
||||||
return;
|
conversation_id: Option<AIConversationId>,
|
||||||
}
|
manual: bool,
|
||||||
|
ctx: &mut ViewContext<Self>,
|
||||||
log::info!("[bedrock] Running AWS login command: {login_command}");
|
) {
|
||||||
|
let exchange_id = conversation_id
|
||||||
// Spawn the login command as an async subprocess so it can open the browser
|
.and_then(|id| Self::thread_start_exchange_id(&id, ctx));
|
||||||
// for SSO authentication. Once it completes, refresh credentials and resume
|
let connection = BedrockAuthConnection::current(ctx);
|
||||||
// the conversation.
|
let conversation_key = conversation_id
|
||||||
|
.map(|id| id.to_string())
|
||||||
|
.unwrap_or_else(|| format!("login-only-{}", self.view_id));
|
||||||
|
let exchange_key = exchange_id.map(|id| id.to_string()).unwrap_or_default();
|
||||||
|
let recovery = BedrockAuthentication::handle(ctx).update(ctx, |authentication, ctx| {
|
||||||
|
authentication.recover(&conversation_key, &exchange_key, manual, ctx)
|
||||||
|
});
|
||||||
let _ = ctx.spawn(
|
let _ = ctx.spawn(
|
||||||
async move {
|
recovery,
|
||||||
let parts: Vec<&str> = login_command.split_whitespace().collect();
|
move |me, result, ctx| {
|
||||||
let Some((cmd, args)) = parts.split_first() else {
|
|
||||||
return Err("Empty login command".to_string());
|
|
||||||
};
|
|
||||||
let result = tokio::process::Command::new(cmd)
|
|
||||||
.args(args)
|
|
||||||
.status()
|
|
||||||
.await;
|
|
||||||
match result {
|
|
||||||
Ok(status) if status.success() => Ok(()),
|
|
||||||
Ok(status) => Err(format!("AWS login command exited with status: {status}")),
|
|
||||||
Err(e) => Err(format!("Failed to spawn AWS login command: {e}")),
|
|
||||||
}
|
|
||||||
},
|
|
||||||
|_me, result, ctx| {
|
|
||||||
match result {
|
match result {
|
||||||
Ok(()) => {
|
Ok(()) => {
|
||||||
log::info!("[bedrock] AWS login completed successfully; reloading credentials before resuming");
|
let Some(conversation_id) = conversation_id else {
|
||||||
|
return;
|
||||||
// Refresh through ApiKeyManager and wait for its completion. A fixed
|
};
|
||||||
// delay can resume the request while the SDK still has stale SSO data.
|
if BedrockAuthConnection::current(ctx) != connection
|
||||||
let refresh = ApiKeyManager::handle(ctx).update(ctx, |manager, ctx| {
|
|| Self::thread_start_exchange_id(&conversation_id, ctx) != exchange_id
|
||||||
crate::ai::aws_credentials::refresh_aws_credentials(manager, ctx)
|
|| BlocklistAIHistoryModel::as_ref(ctx)
|
||||||
});
|
.conversation_status(&conversation_id) != Some(&ConversationStatus::Error)
|
||||||
let _ = ctx.spawn(refresh, |me, result, ctx| {
|
{
|
||||||
match result {
|
return;
|
||||||
Ok(()) => {
|
}
|
||||||
log::info!("[bedrock] AWS credentials refreshed; resuming conversation");
|
log::info!("[bedrock] AWS credentials refreshed; resuming originating conversation");
|
||||||
let conversation_id = if FeatureFlag::AgentView.is_enabled() {
|
me.handle_resume_conversation(&conversation_id, ctx);
|
||||||
me.agent_view_controller
|
|
||||||
.as_ref(ctx)
|
|
||||||
.agent_view_state()
|
|
||||||
.active_conversation_id()
|
|
||||||
} else {
|
|
||||||
BlocklistAIHistoryModel::as_ref(ctx)
|
|
||||||
.last_conversation_id(me.id())
|
|
||||||
};
|
|
||||||
if let Some(conversation_id) = conversation_id {
|
|
||||||
me.handle_resume_conversation(&conversation_id, ctx);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Err(error) => {
|
|
||||||
log::error!(
|
|
||||||
"[bedrock] AWS login succeeded but credential reload failed; not resuming: {error}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
});
|
|
||||||
}
|
}
|
||||||
Err(e) => {
|
Err(error) => {
|
||||||
log::error!("[bedrock] AWS login failed: {e}");
|
log::warn!("[bedrock] AWS authentication recovery stopped: {error}");
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
@@ -20741,36 +20709,12 @@ impl TerminalView {
|
|||||||
ctx.emit(Event::OpenThemeChooser);
|
ctx.emit(Event::OpenThemeChooser);
|
||||||
}
|
}
|
||||||
AIBlockEvent::RunAwsLoginCommand => {
|
AIBlockEvent::RunAwsLoginCommand => {
|
||||||
self.run_aws_login_command(ctx);
|
let conversation_id = block.as_ref(ctx).conversation_id();
|
||||||
|
self.recover_aws_authentication(Some(conversation_id), true, ctx);
|
||||||
}
|
}
|
||||||
AIBlockEvent::RefreshAwsCredentials => {
|
AIBlockEvent::RefreshAwsCredentials => {
|
||||||
let refresh = ApiKeyManager::handle(ctx).update(ctx, |manager, ctx| {
|
let conversation_id = block.as_ref(ctx).conversation_id();
|
||||||
crate::ai::aws_credentials::refresh_aws_credentials(manager, ctx)
|
self.recover_aws_authentication(Some(conversation_id), false, ctx);
|
||||||
});
|
|
||||||
let _ = ctx.spawn(refresh, |me, result, ctx| {
|
|
||||||
match result {
|
|
||||||
Ok(()) => {
|
|
||||||
log::info!("[bedrock] AWS credentials reloaded; resuming conversation");
|
|
||||||
let conversation_id = if FeatureFlag::AgentView.is_enabled() {
|
|
||||||
me.agent_view_controller
|
|
||||||
.as_ref(ctx)
|
|
||||||
.agent_view_state()
|
|
||||||
.active_conversation_id()
|
|
||||||
} else {
|
|
||||||
BlocklistAIHistoryModel::as_ref(ctx).last_conversation_id(me.id())
|
|
||||||
};
|
|
||||||
if let Some(conversation_id) = conversation_id {
|
|
||||||
me.handle_resume_conversation(&conversation_id, ctx);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Err(error) => {
|
|
||||||
log::warn!(
|
|
||||||
"[bedrock] AWS credential reload failed; starting configured login command: {error}"
|
|
||||||
);
|
|
||||||
me.run_aws_login_command(ctx);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
});
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
ctx.notify();
|
ctx.notify();
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
use std::path::PathBuf;
|
use std::sync::Arc;
|
||||||
|
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
|
use futures::future::BoxFuture;
|
||||||
use futures::StreamExt;
|
use futures::StreamExt;
|
||||||
use galaxy_agent_core::{
|
use galaxy_agent_core::{
|
||||||
AgentError, AgentErrorKind, AgentEvent, AgentEventStream, AgentRuntime, ConversationMessage,
|
AgentError, AgentErrorKind, AgentEvent, AgentEventStream, AgentRuntime, ConversationMessage,
|
||||||
@@ -8,35 +9,50 @@ use galaxy_agent_core::{
|
|||||||
TurnRequest, turn_control,
|
TurnRequest, turn_control,
|
||||||
};
|
};
|
||||||
use rig_core::client::CompletionClient;
|
use rig_core::client::CompletionClient;
|
||||||
use rig_core::providers::chatgpt;
|
use rig_core::providers::chatgpt::{self, ChatGPTAuth};
|
||||||
|
|
||||||
use crate::request::build_completion_request;
|
use crate::request::build_completion_request;
|
||||||
use crate::stream::start_model_turn;
|
use crate::stream::start_model_turn;
|
||||||
|
|
||||||
/// Refreshes Galaxy's cached ChatGPT subscription credentials without allowing
|
#[derive(Clone, PartialEq, Eq)]
|
||||||
/// Rig to fall back to its interactive device-code flow.
|
pub struct ChatGPTSubscriptionCredentials {
|
||||||
pub async fn refresh_chatgpt_subscription_credentials() -> Result<(), String> {
|
pub access_token: String,
|
||||||
let client = chatgpt::Client::builder()
|
pub account_id: Option<String>,
|
||||||
.oauth()
|
|
||||||
.allow_device_flow(false)
|
|
||||||
.originator("galaxy")
|
|
||||||
.build()
|
|
||||||
.map_err(|error| error.to_string())?;
|
|
||||||
client.authorize().await.map_err(|error| error.to_string())
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
impl std::fmt::Debug for ChatGPTSubscriptionCredentials {
|
||||||
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||||
|
f.debug_struct("ChatGPTSubscriptionCredentials")
|
||||||
|
.field("access_token", &"<redacted>")
|
||||||
|
.field("account_id", &self.account_id.as_ref().map(|_| "<redacted>"))
|
||||||
|
.finish()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
pub type ChatGPTCredentialResolver = Arc<
|
||||||
|
dyn Fn() -> BoxFuture<'static, Result<ChatGPTSubscriptionCredentials, String>> + Send + Sync,
|
||||||
|
>;
|
||||||
|
|
||||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||||
pub struct ChatGPTSubscriptionRuntimeConfig {
|
pub struct ChatGPTSubscriptionRuntimeConfig {
|
||||||
pub model: String,
|
pub model: String,
|
||||||
pub reasoning_effort: Option<String>,
|
pub reasoning_effort: Option<String>,
|
||||||
pub max_output_tokens: Option<u64>,
|
pub max_output_tokens: Option<u64>,
|
||||||
pub auth_file: Option<PathBuf>,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Clone, Debug)]
|
#[derive(Clone)]
|
||||||
pub struct ChatGPTSubscriptionRuntime {
|
pub struct ChatGPTSubscriptionRuntime {
|
||||||
config: ChatGPTSubscriptionRuntimeConfig,
|
config: ChatGPTSubscriptionRuntimeConfig,
|
||||||
descriptor: RuntimeDescriptor,
|
descriptor: RuntimeDescriptor,
|
||||||
|
credential_resolver: Option<ChatGPTCredentialResolver>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl std::fmt::Debug for ChatGPTSubscriptionRuntime {
|
||||||
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||||
|
f.debug_struct("ChatGPTSubscriptionRuntime")
|
||||||
|
.field("config", &self.config)
|
||||||
|
.field("has_credential_resolver", &self.credential_resolver.is_some())
|
||||||
|
.finish()
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn reasoning_additional_params(reasoning_effort: Option<&str>) -> Option<serde_json::Value> {
|
fn reasoning_additional_params(reasoning_effort: Option<&str>) -> Option<serde_json::Value> {
|
||||||
@@ -56,7 +72,24 @@ impl ChatGPTSubscriptionRuntime {
|
|||||||
kind: RuntimeKind::Provider,
|
kind: RuntimeKind::Provider,
|
||||||
capabilities: RuntimeCapabilities::provider(),
|
capabilities: RuntimeCapabilities::provider(),
|
||||||
};
|
};
|
||||||
Self { config, descriptor }
|
Self { config, descriptor, credential_resolver: None }
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn with_credential_resolver(mut self, resolver: ChatGPTCredentialResolver) -> Self {
|
||||||
|
self.credential_resolver = Some(resolver);
|
||||||
|
self
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn resolve_auth(&self) -> Result<ChatGPTAuth, AgentError> {
|
||||||
|
let resolver = self.credential_resolver.as_ref().ok_or_else(|| {
|
||||||
|
AgentError::new(AgentErrorKind::Authentication, "ChatGPT sign-in required. Reconnect ChatGPT in Providers.")
|
||||||
|
})?;
|
||||||
|
let credentials = resolver().await
|
||||||
|
.map_err(|message| AgentError::new(AgentErrorKind::Authentication, message))?;
|
||||||
|
Ok(ChatGPTAuth::AccessToken {
|
||||||
|
access_token: credentials.access_token,
|
||||||
|
account_id: credentials.account_id,
|
||||||
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Completes a text-only request for callers such as crosscheck reviewers.
|
/// Completes a text-only request for callers such as crosscheck reviewers.
|
||||||
@@ -117,14 +150,11 @@ impl AgentRuntime for ChatGPTSubscriptionRuntime {
|
|||||||
request: TurnRequest,
|
request: TurnRequest,
|
||||||
control: TurnControl,
|
control: TurnControl,
|
||||||
) -> Result<AgentEventStream, AgentError> {
|
) -> Result<AgentEventStream, AgentError> {
|
||||||
let mut builder = chatgpt::Client::builder()
|
let auth = self.resolve_auth().await?;
|
||||||
.oauth()
|
let client = chatgpt::Client::builder()
|
||||||
|
.api_key(auth)
|
||||||
.allow_device_flow(false)
|
.allow_device_flow(false)
|
||||||
.originator("galaxy");
|
.originator("galaxy")
|
||||||
if let Some(auth_file) = &self.config.auth_file {
|
|
||||||
builder = builder.auth_file(auth_file);
|
|
||||||
}
|
|
||||||
let client = builder
|
|
||||||
.build()
|
.build()
|
||||||
.map_err(|error| AgentError::new(AgentErrorKind::Configuration, error.to_string()))?;
|
.map_err(|error| AgentError::new(AgentErrorKind::Configuration, error.to_string()))?;
|
||||||
let model = client.completion_model(&self.config.model);
|
let model = client.completion_model(&self.config.model);
|
||||||
|
|||||||
@@ -1,13 +1,14 @@
|
|||||||
|
use galaxy_agent_core::StopReason;
|
||||||
use rig_core::client::CompletionClient;
|
use rig_core::client::CompletionClient;
|
||||||
use rig_core::completion::{AssistantContent, CompletionModel, Message};
|
use rig_core::completion::{AssistantContent, CompletionModel, Message};
|
||||||
use rig_core::message::{ToolResultContent, UserContent};
|
use rig_core::message::{ToolResultContent, UserContent};
|
||||||
use rig_core::providers::chatgpt::ChatGPTAuth;
|
use rig_core::providers::chatgpt::ChatGPTAuth;
|
||||||
use rig_core::test_utils::RecordingHttpClient;
|
use rig_core::test_utils::{RecordingHttpClient, SequencedStreamingHttpClient};
|
||||||
|
|
||||||
use super::*;
|
use super::*;
|
||||||
|
|
||||||
const COMPLETED_RESPONSE: &str = r#"data: {"type":"response.output_text.delta","delta":"ok"}
|
const COMPLETED_RESPONSE: &str = r#"data: {"type":"response.output_text.delta","delta":"ok"}
|
||||||
data: {"type":"response.completed","response":{"id":"resp_test","object":"response","created_at":1,"status":"completed","error":null,"incomplete_details":null,"instructions":null,"max_output_tokens":null,"model":"gpt-5.3-codex","usage":{"input_tokens":1,"input_tokens_details":{"cached_tokens":0},"output_tokens":1,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":2},"output":[{"type":"message","id":"msg_test","status":"completed","role":"assistant","content":[{"type":"output_text","annotations":[],"text":"ok"}]}],"tools":[]}}
|
data: {"type":"response.completed","sequence_number":54,"response":{"id":"resp_test","object":"response","created_at":1,"status":"completed","error":null,"incomplete_details":null,"instructions":null,"max_output_tokens":null,"model":"gpt-5.3-codex","usage":{"input_tokens":1,"input_tokens_details":{"cached_tokens":0},"output_tokens":1,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":2},"output":[{"type":"message","id":"msg_test","status":"completed","role":"assistant","content":[{"type":"output_text","annotations":[],"text":"ok"}]}],"tools":[]}}
|
||||||
data: [DONE]"#;
|
data: [DONE]"#;
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -16,7 +17,6 @@ fn runtime_descriptor_identifies_chatgpt_subscription() {
|
|||||||
model: "gpt-5.3-codex".to_string(),
|
model: "gpt-5.3-codex".to_string(),
|
||||||
reasoning_effort: Some("high".to_string()),
|
reasoning_effort: Some("high".to_string()),
|
||||||
max_output_tokens: None,
|
max_output_tokens: None,
|
||||||
auth_file: None,
|
|
||||||
});
|
});
|
||||||
|
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
@@ -25,6 +25,43 @@ fn runtime_descriptor_identifies_chatgpt_subscription() {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn runtime_resolves_current_credentials_without_rig_oauth() {
|
||||||
|
let runtime = ChatGPTSubscriptionRuntime::new(ChatGPTSubscriptionRuntimeConfig {
|
||||||
|
model: "test-model".to_owned(),
|
||||||
|
reasoning_effort: None,
|
||||||
|
max_output_tokens: None,
|
||||||
|
}).with_credential_resolver(std::sync::Arc::new(|| {
|
||||||
|
Box::pin(async {
|
||||||
|
Ok(ChatGPTSubscriptionCredentials {
|
||||||
|
access_token: "resolved-token".to_owned(),
|
||||||
|
account_id: Some("resolved-account".to_owned()),
|
||||||
|
})
|
||||||
|
})
|
||||||
|
}));
|
||||||
|
match runtime.resolve_auth().await.unwrap() {
|
||||||
|
ChatGPTAuth::AccessToken { access_token, account_id } => {
|
||||||
|
assert_eq!(access_token, "resolved-token");
|
||||||
|
assert_eq!(account_id.as_deref(), Some("resolved-account"));
|
||||||
|
}
|
||||||
|
ChatGPTAuth::OAuth => panic!("The runtime must not own OAuth refresh"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn missing_or_failed_resolver_requires_sign_in_without_device_flow() {
|
||||||
|
let runtime = ChatGPTSubscriptionRuntime::new(ChatGPTSubscriptionRuntimeConfig {
|
||||||
|
model: "test-model".to_owned(),
|
||||||
|
reasoning_effort: None,
|
||||||
|
max_output_tokens: None,
|
||||||
|
});
|
||||||
|
assert_eq!(runtime.resolve_auth().await.unwrap_err().kind, AgentErrorKind::Authentication);
|
||||||
|
let runtime = runtime.with_credential_resolver(std::sync::Arc::new(|| {
|
||||||
|
Box::pin(async { Err("Sign in again".to_owned()) })
|
||||||
|
}));
|
||||||
|
assert_eq!(runtime.resolve_auth().await.unwrap_err().kind, AgentErrorKind::Authentication);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn reasoning_effort_is_encoded_for_the_responses_request() {
|
fn reasoning_effort_is_encoded_for_the_responses_request() {
|
||||||
let request = TurnRequest::new(
|
let request = TurnRequest::new(
|
||||||
@@ -115,3 +152,71 @@ async fn chatgpt_follow_up_request_preserves_responses_call_ids() {
|
|||||||
assert_eq!(input[0]["call_id"], "call_native_1");
|
assert_eq!(input[0]["call_id"], "call_native_1");
|
||||||
assert_eq!(input[1]["call_id"], "call_native_1");
|
assert_eq!(input[1]["call_id"], "call_native_1");
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn chatgpt_metadata_does_not_interrupt_text_or_completion() {
|
||||||
|
let metadata = r#"data: {"type":"response.metadata","response_id":"resp_test","sequence_number":51,"metadata":{"openai_chatgpt_moderation_metadata":{"prompt":{},"generation":{},"tool_call":{},"tool_response":{}}}}"#;
|
||||||
|
let delta = serde_json::json!({
|
||||||
|
"type": "response.output_text.delta",
|
||||||
|
"item_id": "msg_test",
|
||||||
|
"output_index": 0,
|
||||||
|
"content_index": 0,
|
||||||
|
"sequence_number": 52,
|
||||||
|
"delta": "ok",
|
||||||
|
});
|
||||||
|
let completed = COMPLETED_RESPONSE
|
||||||
|
.lines()
|
||||||
|
.find(|line| line.starts_with("data: {\"type\":\"response.completed\""))
|
||||||
|
.expect("fixture should contain a completed response");
|
||||||
|
let response =
|
||||||
|
format!("{metadata}\n\ndata: {delta}\n\n{metadata}\n\n{completed}\n\ndata: [DONE]\n\n");
|
||||||
|
let client = chatgpt::Client::builder()
|
||||||
|
.api_key(ChatGPTAuth::AccessToken {
|
||||||
|
access_token: "test-token".to_string(),
|
||||||
|
account_id: None,
|
||||||
|
})
|
||||||
|
.http_client(SequencedStreamingHttpClient::new(vec![Ok(response.into())]))
|
||||||
|
.build()
|
||||||
|
.expect("client should build");
|
||||||
|
let request = build_completion_request(
|
||||||
|
TurnRequest::new(
|
||||||
|
"gpt-5.3-codex".to_string(),
|
||||||
|
vec![ConversationMessage {
|
||||||
|
role: MessageRole::User,
|
||||||
|
content: MessageContent::Text("hello".to_string()),
|
||||||
|
}],
|
||||||
|
),
|
||||||
|
None,
|
||||||
|
true,
|
||||||
|
true,
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.expect("request should convert");
|
||||||
|
let (_sender, control) = turn_control();
|
||||||
|
let mut stream = start_model_turn(
|
||||||
|
client.completion_model("gpt-5.3-codex"),
|
||||||
|
request,
|
||||||
|
control,
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.expect("request should reach the mocked provider");
|
||||||
|
let mut text = String::new();
|
||||||
|
let mut usage = None;
|
||||||
|
let mut stop_reason = None;
|
||||||
|
while let Some(event) = stream.next().await {
|
||||||
|
match event.expect("metadata should not fail the turn") {
|
||||||
|
AgentEvent::TextDelta { text: delta } => text.push_str(&delta),
|
||||||
|
AgentEvent::UsageUpdated { usage: updated } => usage = Some(updated),
|
||||||
|
AgentEvent::TurnStopped { reason } => stop_reason = Some(reason),
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
assert_eq!(text, "ok");
|
||||||
|
assert_eq!(
|
||||||
|
usage.expect("completion should report usage").output_tokens,
|
||||||
|
1
|
||||||
|
);
|
||||||
|
assert_eq!(stop_reason, Some(StopReason::Completed));
|
||||||
|
}
|
||||||
|
|||||||
@@ -160,6 +160,13 @@ where
|
|||||||
});
|
});
|
||||||
}
|
}
|
||||||
Ok(StreamedAssistantContent::Unknown(value)) => {
|
Ok(StreamedAssistantContent::Unknown(value)) => {
|
||||||
|
// ChatGPT subscription streams include informational moderation
|
||||||
|
// metadata that Rig does not map to assistant content.
|
||||||
|
if value.get("type").and_then(serde_json::Value::as_str)
|
||||||
|
== Some("response.metadata")
|
||||||
|
{
|
||||||
|
continue;
|
||||||
|
}
|
||||||
if is_keepalive_event(&value) {
|
if is_keepalive_event(&value) {
|
||||||
yield Ok(AgentEvent::KeepAlive);
|
yield Ok(AgentEvent::KeepAlive);
|
||||||
} else {
|
} else {
|
||||||
@@ -512,6 +519,11 @@ fn text_indicates_authentication_failure(text: &str) -> bool {
|
|||||||
"unrecognized client",
|
"unrecognized client",
|
||||||
"security token included in the request is invalid",
|
"security token included in the request is invalid",
|
||||||
"unable to locate credentials",
|
"unable to locate credentials",
|
||||||
|
"unauthorizedexception",
|
||||||
|
"session token not found or invalid",
|
||||||
|
"sso session has expired",
|
||||||
|
"failed to load aws credentials",
|
||||||
|
"bedrock credentials not configured",
|
||||||
"credential should be scoped to a valid region",
|
"credential should be scoped to a valid region",
|
||||||
]
|
]
|
||||||
.iter()
|
.iter()
|
||||||
|
|||||||
@@ -105,6 +105,10 @@ fn flattened_bedrock_credential_errors_are_authentication_failures() {
|
|||||||
r#"{"__type":"ExpiredTokenException","message":"The security token included in the request is expired"}"#,
|
r#"{"__type":"ExpiredTokenException","message":"The security token included in the request is expired"}"#,
|
||||||
"UnrecognizedClientException: The security token included in the request is invalid",
|
"UnrecognizedClientException: The security token included in the request is invalid",
|
||||||
"AccessDeniedException: not authorized to perform bedrock:ConverseStream",
|
"AccessDeniedException: not authorized to perform bedrock:ConverseStream",
|
||||||
|
"UnauthorizedException: Session token not found or invalid",
|
||||||
|
"SSO session has expired or is otherwise invalid",
|
||||||
|
"Failed to load AWS credentials: credentials were not loaded",
|
||||||
|
"Bedrock credentials not configured",
|
||||||
] {
|
] {
|
||||||
let mapped = map_completion_error(CompletionError::ProviderError(message.to_string()));
|
let mapped = map_completion_error(CompletionError::ProviderError(message.to_string()));
|
||||||
assert_eq!(mapped.kind, AgentErrorKind::Authentication);
|
assert_eq!(mapped.kind, AgentErrorKind::Authentication);
|
||||||
@@ -112,3 +116,15 @@ fn flattened_bedrock_credential_errors_are_authentication_failures() {
|
|||||||
assert!(mapped.user_message.is_some());
|
assert!(mapped.user_message.is_some());
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn unrelated_provider_failures_do_not_become_authentication_failures() {
|
||||||
|
for message in [
|
||||||
|
"HTTP client error: connection reset",
|
||||||
|
"ProviderError: temporarily unavailable",
|
||||||
|
"ValidationException: invalid model",
|
||||||
|
] {
|
||||||
|
let mapped = map_completion_error(CompletionError::ProviderError(message.to_owned()));
|
||||||
|
assert_ne!(mapped.kind, AgentErrorKind::Authentication);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user