open_router_migration.rs

  1use credentials_provider::CredentialsProvider;
  2use gpui::App;
  3use util::ResultExt as _;
  4
  5const OPEN_ROUTER_EXTENSION_ID: &str = "openrouter";
  6const OPEN_ROUTER_PROVIDER_ID: &str = "openrouter";
  7const OPEN_ROUTER_DEFAULT_API_URL: &str = "https://openrouter.ai/api/v1";
  8
  9pub fn migrate_open_router_credentials_if_needed(extension_id: &str, cx: &mut App) {
 10    if extension_id != OPEN_ROUTER_EXTENSION_ID {
 11        return;
 12    }
 13
 14    let extension_credential_key = format!(
 15        "extension-llm-{}:{}",
 16        OPEN_ROUTER_EXTENSION_ID, OPEN_ROUTER_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!("OpenRouter extension already has credentials, skipping migration");
 30            return;
 31        }
 32
 33        let old_credential = credentials_provider
 34            .read_credentials(OPEN_ROUTER_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 OpenRouter API key is empty, marking as migrated");
 44                    String::new()
 45                }
 46                Err(_) => {
 47                    log::error!("Failed to decode OpenRouter API key as UTF-8");
 48                    return;
 49                }
 50            },
 51            None => {
 52                log::debug!("No existing OpenRouter 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 OpenRouter API key to OpenRouter 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 OpenRouter API key to extension");
 74            }
 75            Err(err) => {
 76                log::error!("Failed to migrate OpenRouter 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-or-test-key-12345";
 91
 92        cx.write_credentials(OPEN_ROUTER_DEFAULT_API_URL, "Bearer", api_key.as_bytes());
 93
 94        cx.update(|cx| {
 95            migrate_open_router_credentials_if_needed(OPEN_ROUTER_EXTENSION_ID, cx);
 96        });
 97
 98        cx.run_until_parked();
 99
100        let migrated = cx.read_credentials("extension-llm-openrouter:openrouter");
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-or-old-key";
110        let existing_key = "sk-or-existing-key";
111
112        cx.write_credentials(
113            OPEN_ROUTER_DEFAULT_API_URL,
114            "Bearer",
115            old_api_key.as_bytes(),
116        );
117        cx.write_credentials(
118            "extension-llm-openrouter:openrouter",
119            "Bearer",
120            existing_key.as_bytes(),
121        );
122
123        cx.update(|cx| {
124            migrate_open_router_credentials_if_needed(OPEN_ROUTER_EXTENSION_ID, cx);
125        });
126
127        cx.run_until_parked();
128
129        let credentials = cx.read_credentials("extension-llm-openrouter:openrouter");
130        let (_, password) = credentials.unwrap();
131        assert_eq!(
132            String::from_utf8(password).unwrap(),
133            existing_key,
134            "Should not overwrite existing credentials"
135        );
136    }
137
138    #[gpui::test]
139    async fn test_skips_migration_if_empty_marker_exists(cx: &mut TestAppContext) {
140        let old_api_key = "sk-or-old-key";
141
142        // Old credentials exist
143        cx.write_credentials(
144            OPEN_ROUTER_DEFAULT_API_URL,
145            "Bearer",
146            old_api_key.as_bytes(),
147        );
148        // But empty marker already exists (from previous migration attempt)
149        cx.write_credentials("extension-llm-openrouter:openrouter", "Bearer", b"");
150
151        cx.update(|cx| {
152            migrate_open_router_credentials_if_needed(OPEN_ROUTER_EXTENSION_ID, cx);
153        });
154
155        cx.run_until_parked();
156
157        let credentials = cx.read_credentials("extension-llm-openrouter:openrouter");
158        let (_, password) = credentials.unwrap();
159        assert!(
160            password.is_empty(),
161            "Should not overwrite empty marker with old credentials"
162        );
163    }
164
165    #[gpui::test]
166    async fn test_writes_empty_marker_if_no_old_credentials(cx: &mut TestAppContext) {
167        cx.update(|cx| {
168            migrate_open_router_credentials_if_needed(OPEN_ROUTER_EXTENSION_ID, cx);
169        });
170
171        cx.run_until_parked();
172
173        let credentials = cx.read_credentials("extension-llm-openrouter:openrouter");
174        assert!(
175            credentials.is_some(),
176            "Should write empty credentials as migration marker"
177        );
178        let (username, password) = credentials.unwrap();
179        assert_eq!(username, "Bearer");
180        assert!(password.is_empty(), "Password should be empty marker");
181    }
182
183    #[gpui::test]
184    async fn test_skips_migration_for_other_extensions(cx: &mut TestAppContext) {
185        let api_key = "sk-or-test-key";
186
187        cx.write_credentials(OPEN_ROUTER_DEFAULT_API_URL, "Bearer", api_key.as_bytes());
188
189        cx.update(|cx| {
190            migrate_open_router_credentials_if_needed("some-other-extension", cx);
191        });
192
193        cx.run_until_parked();
194
195        let credentials = cx.read_credentials("extension-llm-openrouter:openrouter");
196        assert!(
197            credentials.is_none(),
198            "Should not migrate for other extensions"
199        );
200    }
201}