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