diff --git a/crates/extension_host/src/extension_host.rs b/crates/extension_host/src/extension_host.rs index 591e1347e1518acafbd8d11ede6a3eef98dc1d87..c0b89647a43030538af8d4db6ab80524048a1574 100644 --- a/crates/extension_host/src/extension_host.rs +++ b/crates/extension_host/src/extension_host.rs @@ -2,6 +2,7 @@ mod anthropic_migration; mod capability_granter; mod copilot_migration; pub mod extension_settings; +mod google_ai_migration; pub mod headless_host; pub mod wasm_host; @@ -891,6 +892,7 @@ impl ExtensionStore { // Run extension-specific migrations copilot_migration::migrate_copilot_credentials_if_needed(&extension_id, cx); anthropic_migration::migrate_anthropic_credentials_if_needed(&extension_id, cx); + google_ai_migration::migrate_google_ai_credentials_if_needed(&extension_id, cx); }) .ok(); } diff --git a/crates/extension_host/src/google_ai_migration.rs b/crates/extension_host/src/google_ai_migration.rs new file mode 100644 index 0000000000000000000000000000000000000000..7b946d7d4301c68c01268b1d00ff49f423edd781 --- /dev/null +++ b/crates/extension_host/src/google_ai_migration.rs @@ -0,0 +1,153 @@ +use credentials_provider::CredentialsProvider; +use gpui::App; + +const GOOGLE_AI_EXTENSION_ID: &str = "google-ai"; +const GOOGLE_AI_PROVIDER_ID: &str = "google-ai"; +const GOOGLE_AI_DEFAULT_API_URL: &str = "https://generativelanguage.googleapis.com"; + +pub fn migrate_google_ai_credentials_if_needed(extension_id: &str, cx: &mut App) { + if extension_id != GOOGLE_AI_EXTENSION_ID { + return; + } + + let extension_credential_key = format!( + "extension-llm-{}:{}", + GOOGLE_AI_EXTENSION_ID, GOOGLE_AI_PROVIDER_ID + ); + + let credentials_provider = ::global(cx); + + cx.spawn(async move |cx| { + let existing_credential = credentials_provider + .read_credentials(&extension_credential_key, &cx) + .await + .ok() + .flatten(); + + if existing_credential.is_some() { + log::debug!("Google AI extension already has credentials, skipping migration"); + return; + } + + let old_credential = credentials_provider + .read_credentials(GOOGLE_AI_DEFAULT_API_URL, &cx) + .await + .ok() + .flatten(); + + let api_key = match old_credential { + Some((_, key_bytes)) => match String::from_utf8(key_bytes) { + Ok(key) => key, + Err(_) => { + log::error!("Failed to decode Google AI API key as UTF-8"); + return; + } + }, + None => { + log::debug!("No existing Google AI API key found to migrate"); + return; + } + }; + + log::info!("Migrating existing Google AI API key to Google AI extension"); + + match credentials_provider + .write_credentials(&extension_credential_key, "Bearer", api_key.as_bytes(), &cx) + .await + { + Ok(()) => { + log::info!("Successfully migrated Google AI API key to extension"); + } + Err(err) => { + log::error!("Failed to migrate Google AI API key: {}", err); + } + } + }) + .detach(); +} + +#[cfg(test)] +mod tests { + use super::*; + use gpui::TestAppContext; + + #[gpui::test] + async fn test_migrates_credentials_from_old_location(cx: &mut TestAppContext) { + let api_key = "AIzaSy-test-key-12345"; + + cx.write_credentials(GOOGLE_AI_DEFAULT_API_URL, "Bearer", api_key.as_bytes()); + + cx.update(|cx| { + migrate_google_ai_credentials_if_needed(GOOGLE_AI_EXTENSION_ID, cx); + }); + + cx.run_until_parked(); + + let migrated = cx.read_credentials("extension-llm-google-ai:google-ai"); + assert!(migrated.is_some(), "Credentials should have been migrated"); + let (username, password) = migrated.unwrap(); + assert_eq!(username, "Bearer"); + assert_eq!(String::from_utf8(password).unwrap(), api_key); + } + + #[gpui::test] + async fn test_skips_migration_if_extension_already_has_credentials(cx: &mut TestAppContext) { + let old_api_key = "AIzaSy-old-key"; + let existing_key = "AIzaSy-existing-key"; + + cx.write_credentials(GOOGLE_AI_DEFAULT_API_URL, "Bearer", old_api_key.as_bytes()); + cx.write_credentials( + "extension-llm-google-ai:google-ai", + "Bearer", + existing_key.as_bytes(), + ); + + cx.update(|cx| { + migrate_google_ai_credentials_if_needed(GOOGLE_AI_EXTENSION_ID, cx); + }); + + cx.run_until_parked(); + + let credentials = cx.read_credentials("extension-llm-google-ai:google-ai"); + let (_, password) = credentials.unwrap(); + assert_eq!( + String::from_utf8(password).unwrap(), + existing_key, + "Should not overwrite existing credentials" + ); + } + + #[gpui::test] + async fn test_skips_migration_if_no_old_credentials(cx: &mut TestAppContext) { + cx.update(|cx| { + migrate_google_ai_credentials_if_needed(GOOGLE_AI_EXTENSION_ID, cx); + }); + + cx.run_until_parked(); + + let credentials = cx.read_credentials("extension-llm-google-ai:google-ai"); + assert!( + credentials.is_none(), + "Should not create credentials if none existed" + ); + } + + #[gpui::test] + async fn test_skips_migration_for_other_extensions(cx: &mut TestAppContext) { + let api_key = "AIzaSy-test-key"; + + cx.write_credentials(GOOGLE_AI_DEFAULT_API_URL, "Bearer", api_key.as_bytes()); + + cx.update(|cx| { + migrate_google_ai_credentials_if_needed("some-other-extension", cx); + }); + + cx.run_until_parked(); + + let credentials = cx.read_credentials("extension-llm-google-ai:google-ai"); + assert!( + credentials.is_none(), + "Should not migrate for other extensions" + ); + } +}