open_router_migration.rs

  1use credentials_provider::CredentialsProvider;
  2use gpui::App;
  3
  4const OPEN_ROUTER_EXTENSION_ID: &str = "open-router";
  5const OPEN_ROUTER_PROVIDER_ID: &str = "open-router";
  6const OPEN_ROUTER_DEFAULT_API_URL: &str = "https://openrouter.ai/api/v1";
  7
  8pub fn migrate_open_router_credentials_if_needed(extension_id: &str, cx: &mut App) {
  9    if extension_id != OPEN_ROUTER_EXTENSION_ID {
 10        return;
 11    }
 12
 13    let extension_credential_key = format!(
 14        "extension-llm-{}:{}",
 15        OPEN_ROUTER_EXTENSION_ID, OPEN_ROUTER_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!("OpenRouter extension already has credentials, skipping migration");
 29            return;
 30        }
 31
 32        let old_credential = credentials_provider
 33            .read_credentials(OPEN_ROUTER_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 OpenRouter API key as UTF-8");
 43                    return;
 44                }
 45            },
 46            None => {
 47                log::debug!("No existing OpenRouter API key found to migrate");
 48                return;
 49            }
 50        };
 51
 52        log::info!("Migrating existing OpenRouter API key to OpenRouter 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 OpenRouter API key to extension");
 60            }
 61            Err(err) => {
 62                log::error!("Failed to migrate OpenRouter 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 = "sk-or-test-key-12345";
 77
 78        cx.write_credentials(OPEN_ROUTER_DEFAULT_API_URL, "Bearer", api_key.as_bytes());
 79
 80        cx.update(|cx| {
 81            migrate_open_router_credentials_if_needed(OPEN_ROUTER_EXTENSION_ID, cx);
 82        });
 83
 84        cx.run_until_parked();
 85
 86        let migrated = cx.read_credentials("extension-llm-open-router:open-router");
 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 = "sk-or-old-key";
 96        let existing_key = "sk-or-existing-key";
 97
 98        cx.write_credentials(
 99            OPEN_ROUTER_DEFAULT_API_URL,
100            "Bearer",
101            old_api_key.as_bytes(),
102        );
103        cx.write_credentials(
104            "extension-llm-open-router:open-router",
105            "Bearer",
106            existing_key.as_bytes(),
107        );
108
109        cx.update(|cx| {
110            migrate_open_router_credentials_if_needed(OPEN_ROUTER_EXTENSION_ID, cx);
111        });
112
113        cx.run_until_parked();
114
115        let credentials = cx.read_credentials("extension-llm-open-router:open-router");
116        let (_, password) = credentials.unwrap();
117        assert_eq!(
118            String::from_utf8(password).unwrap(),
119            existing_key,
120            "Should not overwrite existing credentials"
121        );
122    }
123
124    #[gpui::test]
125    async fn test_skips_migration_if_no_old_credentials(cx: &mut TestAppContext) {
126        cx.update(|cx| {
127            migrate_open_router_credentials_if_needed(OPEN_ROUTER_EXTENSION_ID, cx);
128        });
129
130        cx.run_until_parked();
131
132        let credentials = cx.read_credentials("extension-llm-open-router:open-router");
133        assert!(
134            credentials.is_none(),
135            "Should not create credentials if none existed"
136        );
137    }
138
139    #[gpui::test]
140    async fn test_skips_migration_for_other_extensions(cx: &mut TestAppContext) {
141        let api_key = "sk-or-test-key";
142
143        cx.write_credentials(OPEN_ROUTER_DEFAULT_API_URL, "Bearer", api_key.as_bytes());
144
145        cx.update(|cx| {
146            migrate_open_router_credentials_if_needed("some-other-extension", cx);
147        });
148
149        cx.run_until_parked();
150
151        let credentials = cx.read_credentials("extension-llm-open-router:open-router");
152        assert!(
153            credentials.is_none(),
154            "Should not migrate for other extensions"
155        );
156    }
157}