use anyhow::Result;use collections::{HashMap, HashSet};use credentials_provider::CredentialsProvider;use fs::Fs;use futures::Stream;use futures::{FutureExt, StreamExt, future::BoxFuture, stream::BoxStream};use gpui::{App, AsyncApp, Context, Entity, Task, TaskExt};use http_client::{CustomHeaders, HttpClient};use language_model::util::parse_tool_arguments;use language_model::{    ApiKeyState, AuthenticateError, EnvVar, IconOrSvg, InlineDescription, LanguageModel,    LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelId, LanguageModelName,    LanguageModelProvider, LanguageModelProviderId, LanguageModelProviderName,    LanguageModelProviderState, LanguageModelRequest, LanguageModelToolChoice,    LanguageModelToolResultContent, LanguageModelToolUse, MessageContent, ProviderSettingsView,    RateLimiter, Role, StopReason, SubPageProviderSettings, TokenUsage, env_var,};use llama_cpp::{    LLAMA_CPP_API_URL, ModelEntry, Props, get_models, get_props, stream_chat_completion,    stream_model_events,};pub use settings::LlamaCppAvailableModel as AvailableModel;use settings::{Settings, SettingsStore, update_settings_file};use std::pin::Pin;use std::sync::LazyLock;use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};use std::time::Duration;use ui::{    ButtonLike, ButtonLink, ConfiguredApiCard, Divider, List, ListBulletItem, Tooltip, prelude::*,};use ui_input::InputField;use util::ResultExt; use crate::AllLanguageModelSettings; const LLAMA_CPP_DOWNLOAD_URL: &str = "https://llama.app";const LLAMA_CPP_MODELS_URL: &str = "https://huggingface.co/models?library=gguf&sort=trending"; const PROVIDER_ID: LanguageModelProviderId = LanguageModelProviderId::new("llama.cpp");const PROVIDER_NAME: LanguageModelProviderName = LanguageModelProviderName::new("llama.cpp"); const API_KEY_ENV_VAR_NAME: &str = "LLAMACPP_API_KEY";static API_KEY_ENV_VAR: LazyLock<EnvVar> = env_var!(API_KEY_ENV_VAR_NAME); /// How long to wait before reconnecting to `/models/sse` after the stream ends.const MODEL_EVENT_RECONNECT_INTERVAL: Duration = Duration::from_secs(5); /// Context length assumed for an unloaded router model (it can't be probed/// without loading it). Generous so early messages work; re-discovery refines/// it once the model loads.const ASSUMED_UNLOADED_CONTEXT: u64 = 131_072; #[derive(Default, Debug, Clone, PartialEq)]pub struct LlamaCppSettings {    pub api_url: String,    pub auto_discover: bool,    pub available_models: Vec<AvailableModel>,    pub context_window: Option<u64>,    pub custom_headers: CustomHeaders,} pub struct LlamaCppLanguageModelProvider {    http_client: Arc<dyn HttpClient>,    state: Entity<State>,    /// Live capabilities shared with the agent's models (see [`LiveCapabilities`]).    capability_cells: CapabilityCells,    /// Live model-load progress shared with the models (see [`LoadingProgress`]).    loading_progress: LoadingProgress,} pub struct State {    api_key_state: ApiKeyState,    credentials_provider: Arc<dyn CredentialsProvider>,    http_client: Arc<dyn HttpClient>,    fetched_models: Vec<llama_cpp::Model>,    fetch_model_task: Option<Task<Result<()>>>,    /// Router-mode task on `/models/sse`; re-runs discovery as models load/unload.    model_event_task: Option<Task<()>>,    /// Same `Arc` as the provider's; re-discovery keeps these cells in sync.    capability_cells: CapabilityCells,    /// Same `Arc` as the provider's; the event stream updates it as a model loads.    loading_progress: LoadingProgress,} impl State {    fn is_authenticated(&self) -> bool {        !self.fetched_models.is_empty()    }     fn set_api_key(&mut self, api_key: Option<String>, cx: &mut Context<Self>) -> Task<Result<()>> {        let credentials_provider = self.credentials_provider.clone();        let api_url = LlamaCppLanguageModelProvider::api_url(cx);        let task = self.api_key_state.store(            api_url,            api_key,            |this| &mut this.api_key_state,            credentials_provider,            cx,        );         self.fetched_models.clear();        // Drop the event stream so it reconnects with the new key (re-fetch        // below restarts it).        self.model_event_task = None;        write_recover(&self.loading_progress).clear();        cx.spawn(async move |this, cx| {            let result = task.await;            this.update(cx, |this, cx| this.restart_fetch_models_task(cx))                .ok();            result        })    }     fn authenticate(&mut self, cx: &mut Context<Self>) -> Task<Result<(), AuthenticateError>> {        let credentials_provider = self.credentials_provider.clone();        let api_url = LlamaCppLanguageModelProvider::api_url(cx);        let load_key_task = self.api_key_state.load_if_needed(            api_url,            |this| &mut this.api_key_state,            credentials_provider,            cx,        );         if self.is_authenticated() {            return Task::ready(Ok(()));        }         cx.spawn(async move |this, cx| {            match load_key_task.await {                Ok(()) | Err(AuthenticateError::CredentialsNotFound) => {}                Err(error) => {                    log::warn!("failed to load llama.cpp API key: {error}");                }            }            let fetch_models_task = this.update(cx, |this, cx| this.fetch_models(cx))?;            match fetch_models_task.await {                Ok(()) => Ok(()),                Err(err) => {                    // A refused connection means the server isn't running yet, not an error.                    let connection_refused = err.chain().any(|cause| {                        cause                            .downcast_ref::<std::io::Error>()                            .is_some_and(|io_err| {                                io_err.kind() == std::io::ErrorKind::ConnectionRefused                            })                    });                    if connection_refused {                        Err(AuthenticateError::ConnectionRefused)                    } else {                        Err(AuthenticateError::Other(err))                    }                }            }        })    }     fn fetch_models(&mut self, cx: &mut Context<Self>) -> Task<Result<()>> {        let http_client = Arc::clone(&self.http_client);        let settings = LlamaCppLanguageModelProvider::settings(cx);        let api_url = LlamaCppLanguageModelProvider::api_url(cx);        let api_key = self.api_key_state.key(&api_url);        let extra_headers = settings.custom_headers.clone();         cx.spawn(async move |this, cx| {            let entries = get_models(                http_client.as_ref(),                &api_url,                api_key.as_deref(),                &extra_headers,            )            .await?;             let is_router = entries.iter().any(ModelEntry::is_router_entry);             // Models the server reports as loading, used below to prune stale            // progress labels by reconciling against the live listing (a preempted            // load or missed SSE event can skip the terminal event).            let loading_ids: HashSet<String> = entries                .iter()                .filter(|entry| entry.is_loading())                .map(|entry| entry.id.clone())                .collect();             let models: Vec<llama_cpp::Model> = if is_router {                // Router mode: metadata comes from `/v1/models`. We probe                // `/props` only for loaded models so listing never triggers a                // load; unloaded models use the listing's hints and overrides.                let tasks = entries.into_iter().map(|entry| {                    let http_client = Arc::clone(&http_client);                    let api_url = api_url.clone();                    let api_key = api_key.clone();                    let extra_headers = extra_headers.clone();                    async move {                        let props = if entry.is_loaded() {                            get_props(                                http_client.as_ref(),                                &api_url,                                api_key.as_deref(),                                Some(&entry.id),                                &extra_headers,                            )                            .await                            .log_err()                        } else {                            None                        };                        model_from_entry(&entry, props.as_ref())                    }                });                futures::stream::iter(tasks)                    .buffer_unordered(5)                    .collect()                    .await            } else {                // Single-model mode: one `/props` call describes the loaded model.                let props = get_props(                    http_client.as_ref(),                    &api_url,                    api_key.as_deref(),                    None,                    &extra_headers,                )                .await                .log_err();                entries                    .iter()                    .map(|entry| model_from_entry(entry, props.as_ref()))                    .collect()            };             this.update(cx, |this, cx| {                this.fetched_models = models;                let effective = compute_effective_models(                    &this.fetched_models,                    LlamaCppLanguageModelProvider::settings(cx),                );                sync_capability_cells(&this.capability_cells, &effective);                // Drop progress labels for models no longer loading, so a stale                // "Loading …" can't stick after a preempted load or missed event.                write_recover(&this.loading_progress).retain(|id, _| loading_ids.contains(id));                // Router mode loads models on demand: subscribe so capabilities                // self-correct as they load/unload. Start it once (events trigger                // re-discovery, not a re-spawn); single-model mode needs no stream.                if is_router {                    if this.model_event_task.is_none() {                        this.start_model_event_stream(cx);                    }                } else {                    this.model_event_task = None;                }                cx.notify();            })        })    }     /// Subscribes to `/models/sse` and re-runs discovery as models load, unload,    /// or the list changes, so capabilities stay current. Reconnects if the stream    /// drops; on builds without `/models/sse` the refresh is simply skipped.    fn start_model_event_stream(&mut self, cx: &mut Context<Self>) {        let http_client = Arc::clone(&self.http_client);        let api_url = LlamaCppLanguageModelProvider::api_url(cx);        let api_key = self.api_key_state.key(&api_url);        let extra_headers = LlamaCppLanguageModelProvider::settings(cx)            .custom_headers            .clone();         self.model_event_task = Some(cx.spawn(async move |this, cx| {            loop {                match stream_model_events(                    http_client.as_ref(),                    &api_url,                    api_key.as_deref(),                    &extra_headers,                )                .await                {                    Ok(mut events) => {                        while let Some(event) = events.next().await {                            let Some(event) = event.log_err() else {                                continue;                            };                            if let Some(exit_code) = event.load_failure() {                                log::error!(                                    "llama.cpp model {} failed to load (exit code {exit_code})",                                    event.model                                );                            }                            // Loading-progress tick: record it for the selector (no                            // re-discovery). `cx.notify()` drives `ProviderStateChanged`.                            if let Some(progress) = event.load_progress() {                                let label = SharedString::from(progress.progress_label());                                if this                                    .update(cx, |this, cx| {                                        write_recover(&this.loading_progress)                                            .insert(event.model.clone(), label);                                        cx.notify();                                    })                                    .is_err()                                {                                    return;                                }                                continue;                            }                            if !event.changes_model_state() {                                continue;                            }                            // Terminal load/unload (or list change): drop the                            // progress label and re-discover to refresh capabilities.                            if this                                .update(cx, |this, cx| {                                    write_recover(&this.loading_progress).remove(&event.model);                                    this.restart_fetch_models_task(cx);                                })                                .is_err()                            {                                return;                            }                        }                    }                    // Endpoint missing (older build) or connection failed; retry after a backoff.                    Err(error) => {                        log::warn!("llama.cpp model event stream unavailable: {error:#}");                    }                }                 cx.background_executor()                    .timer(MODEL_EVENT_RECONNECT_INTERVAL)                    .await;                if this.update(cx, |_, _| ()).is_err() {                    return;                }            }        }));    }     fn restart_fetch_models_task(&mut self, cx: &mut Context<Self>) {        let task = self.fetch_models(cx);        self.fetch_model_task.replace(task);    }} /// Capabilities that only become known once a router model loads. `LanguageModel`'s/// methods take no `cx`, yet the agent reads them live each turn, so we share/// them through a map that re-discovery updates — an already-selected model/// picks up real values without re-selection. Image support is available from/// `/v1/models` hints, so it stays a field.#[derive(Clone, Copy, Debug, PartialEq)]struct LiveCapabilities {    max_tokens: u64,    supports_tools: bool,    supports_thinking: bool,} impl LiveCapabilities {    fn of(model: &llama_cpp::Model) -> Self {        Self {            max_tokens: model.max_tokens,            supports_tools: model.supports_tools,            supports_thinking: model.supports_thinking,        }    }} /// Live capabilities keyed by model name, shared by the provider and its models.type CapabilityCells = Arc<RwLock<HashMap<String, LiveCapabilities>>>; /// Model name → load-status label (e.g. `"Loading weights 42%"`) while a router/// model loads, shared so the model selector can show progress. Absent once loaded.type LoadingProgress = Arc<RwLock<HashMap<String, SharedString>>>; /// Locks for reading, recovering instead of panicking on a poisoned lock. The/// critical sections are infallible map ops, so poisoning is unreachable anyway.fn read_recover<T>(lock: &RwLock<T>) -> RwLockReadGuard<'_, T> {    lock.read().unwrap_or_else(|poisoned| poisoned.into_inner())} /// Locks for writing; see [`read_recover`].fn write_recover<T>(lock: &RwLock<T>) -> RwLockWriteGuard<'_, T> {    lock.write()        .unwrap_or_else(|poisoned| poisoned.into_inner())} /// The exact set of models `provided_models` exposes: discovery merged with the/// `available_models` and `context_window` overrides. Shared with re-discovery.fn compute_effective_models(    fetched_models: &[llama_cpp::Model],    settings: &LlamaCppSettings,) -> HashMap<String, llama_cpp::Model> {    let mut models: HashMap<String, llama_cpp::Model> = HashMap::default();    if settings.auto_discover {        for model in fetched_models {            let mut model = model.clone();            if let Some(context_window) = settings.context_window {                model.max_tokens = context_window;            }            models.insert(model.name.clone(), model);        }    }    merge_settings_into_models(        &mut models,        &settings.available_models,        settings.context_window,    );    models} /// Updates the shared capability map from the effective models, so a model held/// by an open conversation observes the new values (it reads the map by name).fn sync_capability_cells(cells: &CapabilityCells, effective: &HashMap<String, llama_cpp::Model>) {    let mut cells = write_recover(cells);    for model in effective.values() {        cells.insert(model.name.clone(), LiveCapabilities::of(model));    }} /// Builds a model from a `/v1/models` entry, refined by `/props` when the model/// is loaded. An unloaded router model can't be probed, so we assume optimistic/// capabilities and let re-discovery refine them on load.fn model_from_entry(entry: &ModelEntry, props: Option<&Props>) -> llama_cpp::Model {    let max_tokens = props        .and_then(Props::context_length)        .or_else(|| entry.meta.as_ref().and_then(|meta| meta.n_ctx))        .or_else(|| entry.meta.as_ref().and_then(|meta| meta.n_ctx_train))        .unwrap_or(ASSUMED_UNLOADED_CONTEXT);    // Trust `/props` when present. Without it, assume tools for an unloaded model    // (re-discovery corrects on load) but not for a loaded model whose probe failed.    let supports_tools = match props {        Some(props) => props.supports_tools(),        None => !entry.is_loaded(),    };    let supports_images = props.is_some_and(Props::supports_images) || entry.supports_images_hint();    let supports_thinking = props.is_some_and(Props::supports_thinking);     llama_cpp::Model::new(        &entry.id,        Some(&display_name_for(&entry.id)),        Some(max_tokens),        supports_tools,        supports_images,        supports_thinking,    )} /// Friendly display name from a model id, which is often a `.gguf` file path.fn display_name_for(id: &str) -> String {    let base = id.rsplit(['/', '\\']).next().unwrap_or(id);    base.strip_suffix(".gguf").unwrap_or(base).to_string()} fn telemetry_id_for(id: &str) -> String {    format!("{PROVIDER_ID}/{}", display_name_for(id))} impl LlamaCppLanguageModelProvider {    pub fn new(        http_client: Arc<dyn HttpClient>,        credentials_provider: Arc<dyn CredentialsProvider>,        cx: &mut App,    ) -> Self {        let capability_cells: CapabilityCells = Arc::new(RwLock::new(HashMap::default()));        let loading_progress: LoadingProgress = Arc::new(RwLock::new(HashMap::default()));        let this = Self {            http_client: http_client.clone(),            capability_cells: capability_cells.clone(),            loading_progress: loading_progress.clone(),            state: cx.new(|cx| {                cx.observe_global::<SettingsStore>({                    let mut last_settings = LlamaCppLanguageModelProvider::settings(cx).clone();                    move |this: &mut State, cx| {                        let current_settings = LlamaCppLanguageModelProvider::settings(cx);                        let settings_changed = current_settings != &last_settings;                        if settings_changed {                            let url_changed = last_settings.api_url != current_settings.api_url;                            last_settings = current_settings.clone();                            if url_changed {                                let credentials_provider = this.credentials_provider.clone();                                let api_url = Self::api_url(cx);                                this.api_key_state.handle_url_change(                                    api_url,                                    |this| &mut this.api_key_state,                                    credentials_provider,                                    cx,                                );                                this.fetched_models.clear();                                // Drop the event stream so it reconnects against                                // the new URL (re-auth below restarts it).                                this.model_event_task = None;                                write_recover(&this.loading_progress).clear();                                this.authenticate(cx).detach();                            }                            cx.notify();                        }                    }                })                .detach();                 State {                    http_client,                    fetched_models: Default::default(),                    fetch_model_task: None,

Lines 1–500 of 2120. Load more