extension.rs

 1use collections::HashMap;
 2use extension::{
 3    ExtensionHostProxy, ExtensionLanguageModelProviderProxy, LanguageModelProviderRegistration,
 4};
 5use gpui::{App, Entity};
 6use language_model::{LanguageModelProviderId, LanguageModelRegistry};
 7use std::sync::{Arc, LazyLock};
 8
 9/// Maps built-in provider IDs to their corresponding extension IDs.
10/// When an extension with this ID is installed, the built-in provider should be hidden.
11static BUILTIN_TO_EXTENSION_MAP: LazyLock<HashMap<&'static str, &'static str>> =
12    LazyLock::new(|| {
13        let mut map = HashMap::default();
14        map.insert("anthropic", "anthropic");
15        map.insert("openai", "openai");
16        map.insert("google", "google-ai");
17        map.insert("openrouter", "openrouter");
18        map.insert("copilot_chat", "copilot-chat");
19        map
20    });
21
22/// Returns the extension ID that should hide the given built-in provider.
23pub fn extension_for_builtin_provider(provider_id: &str) -> Option<&'static str> {
24    BUILTIN_TO_EXTENSION_MAP.get(provider_id).copied()
25}
26
27/// Proxy that registers extension language model providers with the LanguageModelRegistry.
28pub struct LanguageModelProviderRegistryProxy {
29    registry: Entity<LanguageModelRegistry>,
30}
31
32impl LanguageModelProviderRegistryProxy {
33    pub fn new(registry: Entity<LanguageModelRegistry>) -> Self {
34        Self { registry }
35    }
36}
37
38impl ExtensionLanguageModelProviderProxy for LanguageModelProviderRegistryProxy {
39    fn register_language_model_provider(
40        &self,
41        _provider_id: Arc<str>,
42        register_fn: LanguageModelProviderRegistration,
43        cx: &mut App,
44    ) {
45        register_fn(cx);
46    }
47
48    fn unregister_language_model_provider(&self, provider_id: Arc<str>, cx: &mut App) {
49        self.registry.update(cx, |registry, cx| {
50            registry.unregister_provider(LanguageModelProviderId::from(provider_id), cx);
51        });
52    }
53}
54
55/// Initialize the extension language model provider proxy.
56/// This must be called BEFORE extension_host::init to ensure the proxy is available
57/// when extensions try to register their language model providers.
58pub fn init_proxy(cx: &mut App) {
59    let proxy = ExtensionHostProxy::default_global(cx);
60    let registry = LanguageModelRegistry::global(cx);
61
62    registry.update(cx, |registry, _cx| {
63        registry.set_builtin_provider_hiding_fn(Box::new(extension_for_builtin_provider));
64    });
65
66    proxy.register_language_model_provider_proxy(LanguageModelProviderRegistryProxy::new(registry));
67}