feat: region expansion

crates/amzn-toolkit-telemetry/src/config/endpoint.rs

@@ -5,7 +5,6 @@ pub use ::aws_smithy_runtime_api::client::endpoint::{
 };
 pub use ::aws_smithy_types::endpoint::Endpoint;
 
-///
 #[cfg(test)]
 mod test {}
 

crates/fig_api_client/src/clients/client.rs

@@ -37,6 +37,7 @@ use crate::model::{
     RecommendationsInput,
     RecommendationsOutput,
 };
+use crate::profile::Profile;
 use crate::{
     Customization,
     Endpoint,
@@ -207,6 +208,31 @@ impl Client {
             inner::Inner::Mock => Ok(()),
         }
     }
+
+    pub async fn list_available_profiles(&self) -> Result<Vec<Profile>, Error> {
+        match &self.0 {
+            inner::Inner::Codewhisperer(client) => {
+                let mut profiles = vec![];
+                while let Some(profiles_output) = client.list_available_profiles().into_paginator().send().next().await
+                {
+                    profiles.extend(profiles_output?.profiles().iter().cloned().map(Profile::from));
+                }
+
+                Ok(profiles)
+            },
+            inner::Inner::Consolas(_) => Err(Error::UnsupportedConsolas("list_available_profiles")),
+            inner::Inner::Mock => Ok(vec![
+                Profile {
+                    arn: "my:arn:1".to_owned(),
+                    profile_name: "MyProfile".to_owned(),
+                },
+                Profile {
+                    arn: "my:arn:2".to_owned(),
+                    profile_name: "MyOtherProfile".to_owned(),
+                },
+            ]),
+        }
+    }
 }
 
 async fn codewhisperer_generate_recommendation_inner(

crates/fig_api_client/src/clients/mod.rs

@@ -1,5 +1,5 @@
 mod client;
-mod shared;
+pub(crate) mod shared;
 mod streaming_client;
 
 pub use client::{

crates/fig_api_client/src/consts.rs

@@ -7,6 +7,10 @@ pub const PROD_CODEWHISPERER_ENDPOINT_REGION: Region = Region::from_static("us-e
 pub const PROD_Q_ENDPOINT_URL: &str = "https://q.us-east-1.amazonaws.com";
 pub const PROD_Q_ENDPOINT_REGION: Region = Region::from_static("us-east-1");
 
+// FRA endpoint constants
+pub const PROD_CODEWHISPERER_FRA_ENDPOINT_URL: &str = "https://codewhisperer.eu-central-1.amazonaws.com";
+pub const PROD_CODEWHISPERER_FRA_ENDPOINT_REGION: Region = Region::from_static("eu-central-1");
+
 // Opt out constants
 pub const SHARE_CODEWHISPERER_CONTENT_SETTINGS_KEY: &str = "codeWhisperer.shareCodeWhispererContentWithAWS";
 pub const X_AMZN_CODEWHISPERER_OPT_OUT_HEADER: &str = "x-amzn-codewhisperer-optout";

crates/fig_api_client/src/endpoints.rs

@@ -6,41 +6,56 @@ use serde_json::Value;
 use crate::consts::{
     PROD_CODEWHISPERER_ENDPOINT_REGION,
     PROD_CODEWHISPERER_ENDPOINT_URL,
+    PROD_CODEWHISPERER_FRA_ENDPOINT_REGION,
+    PROD_CODEWHISPERER_FRA_ENDPOINT_URL,
     PROD_Q_ENDPOINT_REGION,
     PROD_Q_ENDPOINT_URL,
 };
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct Endpoint {
-    url: Cow<'static, str>,
-    region: Region,
+    pub url: Cow<'static, str>,
+    pub region: Region,
 }
 
 impl Endpoint {
-    const PROD_CODEWHISPERER: Self = Self {
+    pub const CODEWHISPERER_ENDPOINTS: [Self; 2] = [Self::DEFAULT_ENDPOINT, Self {
+        url: Cow::Borrowed(PROD_CODEWHISPERER_FRA_ENDPOINT_URL),
+        region: PROD_CODEWHISPERER_FRA_ENDPOINT_REGION,
+    }];
+    pub const DEFAULT_ENDPOINT: Self = Self {
         url: Cow::Borrowed(PROD_CODEWHISPERER_ENDPOINT_URL),
         region: PROD_CODEWHISPERER_ENDPOINT_REGION,
     };
-    const PROD_Q: Self = Self {
+    pub const PROD_Q: Self = Self {
         url: Cow::Borrowed(PROD_Q_ENDPOINT_URL),
         region: PROD_Q_ENDPOINT_REGION,
     };
 
     pub fn load_codewhisperer() -> Self {
-        match fig_settings::settings::get_value("api.codewhisperer.service") {
-            Ok(Some(Value::Object(o))) => {
-                let endpoint = o.get("endpoint").and_then(|v| v.as_str());
-                let region = o.get("region").and_then(|v| v.as_str());
+        let (endpoint, region) =
+            if let Ok(Some(Value::Object(o))) = fig_settings::settings::get_value("api.codewhisperer.service") {
+                // The following branch is evaluated in case the user has set their own endpoint.
+                (
+                    o.get("endpoint").and_then(|v| v.as_str()).map(|v| v.to_owned()),
+                    o.get("region").and_then(|v| v.as_str()).map(|v| v.to_owned()),
+                )
+            } else if let Ok(Some(Value::Object(o))) = fig_settings::state::get_value("api.codewhisperer.service") {
+                // The following branch is evaluated in the case of user profile being set.
+                (
+                    o.get("endpoint").and_then(|v| v.as_str()).map(|v| v.to_owned()),
+                    o.get("region").and_then(|v| v.as_str()).map(|v| v.to_owned()),
+                )
+            } else {
+                (None, None)
+            };
 
-                match (endpoint, region) {
-                    (Some(endpoint), Some(region)) => Self {
-                        url: endpoint.to_owned().into(),
-                        region: Region::new(region.to_owned()),
-                    },
-                    _ => Endpoint::PROD_CODEWHISPERER,
-                }
+        match (endpoint, region) {
+            (Some(endpoint), Some(region)) => Self {
+                url: endpoint.clone().into(),
+                region: Region::new(region.clone()),
             },
-            _ => Endpoint::PROD_CODEWHISPERER,
+            _ => Endpoint::DEFAULT_ENDPOINT,
         }
     }
 
@@ -82,7 +97,7 @@ mod tests {
         let _ = Endpoint::load_codewhisperer();
         let _ = Endpoint::load_q();
 
-        let prod = Endpoint::PROD_CODEWHISPERER;
+        let prod = &Endpoint::CODEWHISPERER_ENDPOINTS[0];
         Url::parse(prod.url()).unwrap();
         assert_eq!(prod.region(), &PROD_CODEWHISPERER_ENDPOINT_REGION);
 

crates/fig_api_client/src/error.rs

@@ -1,8 +1,7 @@
 use amzn_codewhisperer_client::operation::generate_completions::GenerateCompletionsError;
 use amzn_codewhisperer_client::operation::list_available_customizations::ListAvailableCustomizationsError;
+use amzn_codewhisperer_client::operation::list_available_profiles::ListAvailableProfilesError;
 pub use amzn_codewhisperer_streaming_client::operation::generate_assistant_response::GenerateAssistantResponseError;
-// use amzn_codewhisperer_streaming_client::operation::send_message::SendMessageError as
-// CodewhispererSendMessageError;
 use amzn_codewhisperer_streaming_client::types::error::ChatResponseStreamError as CodewhispererChatResponseStreamError;
 use amzn_consolas_client::operation::generate_recommendations::GenerateRecommendationsError;
 use amzn_consolas_client::operation::list_customizations::ListCustomizationsError;
@@ -61,6 +60,9 @@ pub enum Error {
 
     #[error("unsupported action by consolas: {0}")]
     UnsupportedConsolas(&'static str),
+
+    #[error(transparent)]
+    ListAvailableProfilesError(#[from] SdkError<ListAvailableProfilesError, HttpResponse>),
 }
 
 impl Error {
@@ -75,6 +77,7 @@ impl Error {
                 e.as_service_error().is_some_and(|e| e.is_throttling_error())
             },
             Error::QDeveloperSendMessage(e) => e.as_service_error().is_some_and(|e| e.is_throttling_error()),
+            Error::ListAvailableProfilesError(e) => e.as_service_error().is_some_and(|e| e.is_throttling_error()),
             Error::CodewhispererChatResponseStream(_)
             | Error::QDeveloperChatResponseStream(_)
             | Error::SmithyBuild(_)
@@ -94,6 +97,7 @@ impl Error {
             Error::CodewhispererGenerateAssistantResponse(e) => e.as_service_error().is_some(),
             Error::QDeveloperSendMessage(e) => e.as_service_error().is_some(),
             Error::ContextWindowOverflow => true,
+            Error::ListAvailableProfilesError(e) => e.as_service_error().is_some(),
             Error::CodewhispererChatResponseStream(_)
             | Error::QDeveloperChatResponseStream(_)
             | Error::SmithyBuild(_)

crates/fig_api_client/src/lib.rs

@@ -6,6 +6,7 @@ mod endpoints;
 pub mod error;
 pub(crate) mod interceptor;
 pub mod model;
+pub mod profile;
 
 pub use clients::{
     Client,
@@ -14,3 +15,4 @@ pub use clients::{
 pub use customization::Customization;
 pub use endpoints::Endpoint;
 pub use error::Error;
+pub use profile::list_available_profiles;

crates/fig_api_client/src/profile.rs

@@ -0,0 +1,35 @@
+use serde::{
+    Deserialize,
+    Serialize,
+};
+
+use crate::Client;
+use crate::endpoints::Endpoint;
+
+#[derive(Debug, Deserialize, Serialize)]
+pub struct Profile {
+    pub arn: String,
+    pub profile_name: String,
+}
+
+impl From<amzn_codewhisperer_client::types::Profile> for Profile {
+    fn from(profile: amzn_codewhisperer_client::types::Profile) -> Self {
+        Self {
+            arn: profile.arn,
+            profile_name: profile.profile_name,
+        }
+    }
+}
+
+pub async fn list_available_profiles() -> Vec<Profile> {
+    let mut profiles = vec![];
+    for endpoint in Endpoint::CODEWHISPERER_ENDPOINTS {
+        let client = Client::new_codewhisperer_client(&endpoint).await;
+        match client.list_available_profiles().await {
+            Ok(mut p) => profiles.append(&mut p),
+            Err(e) => tracing::warn!("Failed to list profiles from endpoint {:?}: {:?}", endpoint, e),
+        }
+    }
+
+    profiles
+}

crates/fig_auth/src/profile.rs

@@ -0,0 +1,26 @@
+use serde::{
+    Deserialize,
+    Serialize,
+};
+
+#[derive(Debug, Deserialize, Serialize)]
+pub struct Profile {
+    pub arn: String,
+    pub profile_name: String,
+}
+
+impl From<amzn_codewhisperer_client::types::Profile> for Profile {
+    fn from(profile: amzn_codewhisperer_client::types::Profile) -> Self {
+        Self {
+            arn: profile.arn,
+            profile_name: profile.profile_name,
+        }
+    }
+}
+
+pub fn set_profile(profile: Profile) -> Result<(), Error> {
+    Ok(fig_settings::state::set_value(
+        "api.codewhisperer.service",
+        serde_json::to_string(&profile)?,
+    )?)
+}

crates/fig_desktop_api/src/handler.rs

@@ -182,13 +182,15 @@ where
                 HistoryQueryRequest,
                 InsertTextRequest,
                 InstallRequest,
+                ListAvailableProfilesRequest,
                 NotificationRequest,
                 OnboardingRequest,
                 OpenInExternalApplicationRequest,
                 PingRequest,
                 PositionWindowRequest,
                 ReadFileRequest,
                 RunProcessRequest,
+                SetProfileRequest,
                 TelemetryPageRequest,
                 TelemetryTrackRequest,
                 UpdateApplicationPropertiesRequest,
@@ -265,6 +267,8 @@ where
                 CheckForUpdatesRequest(request) => update::check_for_updates(request).await,
                 GetPlatformInfoRequest(request) => platform::get_platform_info(request, &ctx).await,
                 UserLogoutRequest(request) => event_handler.user_logout(request!(request)).await,
+                ListAvailableProfilesRequest(request) => profile::list_available_profiles(request).await,
+                SetProfileRequest(request) => profile::set_profile(request),
             }
         },
         None => {

crates/fig_desktop_api/src/requests/mod.rs

@@ -9,6 +9,7 @@ pub mod history;
 pub mod install;
 pub mod other;
 pub mod platform;
+pub mod profile;
 pub mod settings;
 pub mod state;
 pub mod telemetry;

crates/fig_desktop_api/src/requests/profile.rs

@@ -0,0 +1,50 @@
+use fig_api_client::profile::Profile;
+use fig_proto::fig::{
+    ListAvailableProfilesRequest,
+    ListAvailableProfilesResponse,
+    SetProfileRequest,
+};
+
+use super::{
+    RequestResult,
+    RequestResultImpl,
+    ServerOriginatedSubMessage,
+};
+
+pub fn set_profile(request: SetProfileRequest) -> RequestResult {
+    let Some(profile) = request.profile else {
+        return RequestResult::error("Profile was not provided.");
+    };
+
+    let profile = Profile {
+        arn: profile.arn,
+        profile_name: profile.profile_name,
+    };
+
+    let profile = match serde_json::to_string(&profile) {
+        Ok(profile) => profile,
+        Err(err) => return RequestResult::error(err.to_string()),
+    };
+
+    if let Err(err) = fig_settings::state::set_value("api.codewhisperer.profile", profile) {
+        return RequestResult::error(err.to_string());
+    }
+
+    RequestResult::success()
+}
+
+pub async fn list_available_profiles(_request: ListAvailableProfilesRequest) -> RequestResult {
+    Ok(
+        ServerOriginatedSubMessage::ListAvailableProfilesResponse(ListAvailableProfilesResponse {
+            profiles: fig_api_client::profile::list_available_profiles()
+                .await
+                .iter()
+                .map(|profile| fig_proto::fig::Profile {
+                    arn: profile.arn.clone(),
+                    profile_name: profile.profile_name.clone(),
+                })
+                .collect(),
+        })
+        .into(),
+    )
+}

crates/q_cli/src/cli/chat/skim_integration.rs

@@ -329,7 +329,7 @@ mod tests {
     #[test]
     fn test_hardcoded_commands_in_commands_array() {
         // Get the set of available commands from prompt.rs
-        let available_commands: HashSet<String> = get_available_commands().iter().map(|cmd| cmd.clone()).collect();
+        let available_commands: HashSet<String> = get_available_commands().iter().cloned().collect();
 
         // List of hardcoded commands used in select_command
         let hardcoded_commands = vec![

proto/fig.proto

@@ -44,6 +44,8 @@ message ClientOriginatedMessage {
     AuthFinishPkceAuthorizationRequest auth_finish_pkce_authorization_request = 150;
     AuthCancelPkceAuthorizationRequest auth_cancel_pkce_authorization_request = 151;
     GetPlatformInfoRequest get_platform_info_request = 152;
+    ListAvailableProfilesRequest list_available_profiles_request = 155;
+    SetProfileRequest set_profile_request = 156;
   }
 
   reserved 115;
@@ -75,6 +77,7 @@ message ServerOriginatedMessage {
     AuthFinishPkceAuthorizationResponse auth_finish_pkce_authorization_response = 123;
     AuthCancelPkceAuthorizationResponse auth_cancel_pkce_authorization_response = 124;
     GetPlatformInfoResponse get_platform_info_response = 125;
+    ListAvailableProfilesResponse list_available_profiles_response = 155;
 
     Notification notification = 1000;
   }
@@ -378,6 +381,21 @@ message GetPlatformInfoResponse {
   optional AppBundleType app_bundle_type = 4;
 }
 
+message Profile {
+  string arn = 1;
+  string profileName = 2;
+}
+
+message ListAvailableProfilesRequest {}
+
+message ListAvailableProfilesResponse {
+  repeated Profile profiles = 1;
+}
+
+message SetProfileRequest {
+  Profile profile = 1;
+}
+
 /// User
 
 // todo

rust-toolchain.toml

@@ -1,5 +1,5 @@
 [toolchain]
-channel = "1.84.0"
+channel = "1.85.0"
 profile = "minimal"
 components = ["rustfmt", "clippy"]
 targets = [