1use credentials_provider::CredentialsProvider;
2use gpui::App;
3
4const GOOGLE_AI_EXTENSION_ID: &str = "google-ai";
5const GOOGLE_AI_PROVIDER_ID: &str = "google-ai";
6const GOOGLE_AI_DEFAULT_API_URL: &str = "https://generativelanguage.googleapis.com";
7
8pub fn migrate_google_ai_credentials_if_needed(extension_id: &str, cx: &mut App) {
9 if extension_id != GOOGLE_AI_EXTENSION_ID {
10 return;
11 }
12
13 let extension_credential_key = format!(
14 "extension-llm-{}:{}",
15 GOOGLE_AI_EXTENSION_ID, GOOGLE_AI_PROVIDER_ID
16 );
17
18 let credentials_provider = <dyn CredentialsProvider>::global(cx);
19
20 cx.spawn(async move |cx| {
21 let existing_credential = credentials_provider
22 .read_credentials(&extension_credential_key, &cx)
23 .await
24 .ok()
25 .flatten();
26
27 if existing_credential.is_some() {
28 log::debug!("Google AI extension already has credentials, skipping migration");
29 return;
30 }
31
32 let old_credential = credentials_provider
33 .read_credentials(GOOGLE_AI_DEFAULT_API_URL, &cx)
34 .await
35 .ok()
36 .flatten();
37
38 let api_key = match old_credential {
39 Some((_, key_bytes)) => match String::from_utf8(key_bytes) {
40 Ok(key) => key,
41 Err(_) => {
42 log::error!("Failed to decode Google AI API key as UTF-8");
43 return;
44 }
45 },
46 None => {
47 log::debug!("No existing Google AI API key found to migrate");
48 return;
49 }
50 };
51
52 log::info!("Migrating existing Google AI API key to Google AI extension");
53
54 match credentials_provider
55 .write_credentials(&extension_credential_key, "Bearer", api_key.as_bytes(), &cx)
56 .await
57 {
58 Ok(()) => {
59 log::info!("Successfully migrated Google AI API key to extension");
60 }
61 Err(err) => {
62 log::error!("Failed to migrate Google AI API key: {}", err);
63 }
64 }
65 })
66 .detach();
67}
68
69#[cfg(test)]
70mod tests {
71 use super::*;
72 use gpui::TestAppContext;
73
74 #[gpui::test]
75 async fn test_migrates_credentials_from_old_location(cx: &mut TestAppContext) {
76 let api_key = "AIzaSy-test-key-12345";
77
78 cx.write_credentials(GOOGLE_AI_DEFAULT_API_URL, "Bearer", api_key.as_bytes());
79
80 cx.update(|cx| {
81 migrate_google_ai_credentials_if_needed(GOOGLE_AI_EXTENSION_ID, cx);
82 });
83
84 cx.run_until_parked();
85
86 let migrated = cx.read_credentials("extension-llm-google-ai:google-ai");
87 assert!(migrated.is_some(), "Credentials should have been migrated");
88 let (username, password) = migrated.unwrap();
89 assert_eq!(username, "Bearer");
90 assert_eq!(String::from_utf8(password).unwrap(), api_key);
91 }
92
93 #[gpui::test]
94 async fn test_skips_migration_if_extension_already_has_credentials(cx: &mut TestAppContext) {
95 let old_api_key = "AIzaSy-old-key";
96 let existing_key = "AIzaSy-existing-key";
97
98 cx.write_credentials(GOOGLE_AI_DEFAULT_API_URL, "Bearer", old_api_key.as_bytes());
99 cx.write_credentials(
100 "extension-llm-google-ai:google-ai",
101 "Bearer",
102 existing_key.as_bytes(),
103 );
104
105 cx.update(|cx| {
106 migrate_google_ai_credentials_if_needed(GOOGLE_AI_EXTENSION_ID, cx);
107 });
108
109 cx.run_until_parked();
110
111 let credentials = cx.read_credentials("extension-llm-google-ai:google-ai");
112 let (_, password) = credentials.unwrap();
113 assert_eq!(
114 String::from_utf8(password).unwrap(),
115 existing_key,
116 "Should not overwrite existing credentials"
117 );
118 }
119
120 #[gpui::test]
121 async fn test_skips_migration_if_no_old_credentials(cx: &mut TestAppContext) {
122 cx.update(|cx| {
123 migrate_google_ai_credentials_if_needed(GOOGLE_AI_EXTENSION_ID, cx);
124 });
125
126 cx.run_until_parked();
127
128 let credentials = cx.read_credentials("extension-llm-google-ai:google-ai");
129 assert!(
130 credentials.is_none(),
131 "Should not create credentials if none existed"
132 );
133 }
134
135 #[gpui::test]
136 async fn test_skips_migration_for_other_extensions(cx: &mut TestAppContext) {
137 let api_key = "AIzaSy-test-key";
138
139 cx.write_credentials(GOOGLE_AI_DEFAULT_API_URL, "Bearer", api_key.as_bytes());
140
141 cx.update(|cx| {
142 migrate_google_ai_credentials_if_needed("some-other-extension", cx);
143 });
144
145 cx.run_until_parked();
146
147 let credentials = cx.read_credentials("extension-llm-google-ai:google-ai");
148 assert!(
149 credentials.is_none(),
150 "Should not migrate for other extensions"
151 );
152 }
153}