vtcode_llm/providers/
local_server.rs

1use std::collections::HashMap;
2use std::process::Stdio;
3use std::sync::{LazyLock, Mutex};
4use std::time::Duration;
5
6use anyhow::{Context, Result};
7use serde::Deserialize;
8use tokio::process::Child;
9
10use crate::process_env::{sanitize_std_command_environment, sanitize_tokio_command_environment};
11use vtcode_config::constants::{env_vars, urls};
12
13use super::common::override_base_url;
14
15const PROBE_TIMEOUT: Duration = Duration::from_secs(5);
16
17// ---------------------------------------------------------------------------
18// Public types
19// ---------------------------------------------------------------------------
20
21#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
22pub enum LocalProvider {
23    Ollama,
24    LmStudio,
25    LlamaCpp,
26}
27
28impl LocalProvider {
29    pub fn key(self) -> &'static str {
30        match self {
31            Self::Ollama => "ollama",
32            Self::LmStudio => "lmstudio",
33            Self::LlamaCpp => "llamacpp",
34        }
35    }
36
37    pub fn display_name(self) -> &'static str {
38        match self {
39            Self::Ollama => "Ollama",
40            Self::LmStudio => "LM Studio",
41            Self::LlamaCpp => "llama.cpp",
42        }
43    }
44
45    pub fn from_key(key: &str) -> Option<Self> {
46        match key {
47            "ollama" => Some(Self::Ollama),
48            "lmstudio" | "lm-studio" => Some(Self::LmStudio),
49            "llamacpp" | "llama.cpp" | "llama-cpp" => Some(Self::LlamaCpp),
50            _ => None,
51        }
52    }
53
54    pub fn all() -> &'static [LocalProvider] {
55        &[Self::Ollama, Self::LmStudio, Self::LlamaCpp]
56    }
57
58    fn default_port(self) -> u16 {
59        match self {
60            Self::Ollama => 11434,
61            Self::LmStudio => 1234,
62            Self::LlamaCpp => 8080,
63        }
64    }
65
66    fn default_base_url(self) -> &'static str {
67        match self {
68            Self::Ollama => urls::OLLAMA_API_BASE,
69            Self::LmStudio => urls::LMSTUDIO_API_BASE,
70            Self::LlamaCpp => urls::LLAMACPP_API_BASE,
71        }
72    }
73
74    fn base_url_env(self) -> &'static str {
75        match self {
76            Self::Ollama => env_vars::OLLAMA_BASE_URL,
77            Self::LmStudio => env_vars::LMSTUDIO_BASE_URL,
78            Self::LlamaCpp => env_vars::LLAMACPP_BASE_URL,
79        }
80    }
81
82    fn base_url(self) -> String {
83        // Shared trim/empty-filter chain: no explicit override exists at
84        // this call site, so resolution is env-first and default-fallback.
85        override_base_url(self.default_base_url(), None, Some(self.base_url_env()))
86    }
87
88    fn host_root(self) -> String {
89        let base = self.base_url();
90        strip_path_suffix(&base)
91    }
92}
93
94#[derive(Debug, Clone)]
95pub struct LocalServerStatus {
96    pub provider: LocalProvider,
97    pub running: bool,
98    pub endpoint: String,
99    pub available_models: Vec<String>,
100    pub running_models: Vec<String>,
101    pub version: Option<String>,
102    pub error: Option<String>,
103}
104
105impl LocalServerStatus {
106    fn not_running(provider: LocalProvider, reason: impl Into<String>) -> Self {
107        Self {
108            provider,
109            running: false,
110            endpoint: provider.base_url(),
111            available_models: Vec::new(),
112            running_models: Vec::new(),
113            version: None,
114            error: Some(reason.into()),
115        }
116    }
117}
118
119#[derive(Debug, Clone)]
120pub struct LocalServerCapabilities {
121    can_start: bool,
122    can_stop: bool,
123    binary_found: bool,
124    binary_name: &'static str,
125    binary_path: Option<String>,
126}
127
128#[derive(Debug, Clone)]
129pub struct EnvVarInfo {
130    pub name: &'static str,
131    pub current_value: Option<String>,
132    pub description: &'static str,
133}
134
135// ---------------------------------------------------------------------------
136// Managed process tracking (for Ollama and llama.cpp)
137// ---------------------------------------------------------------------------
138
139struct ManagedProcess {
140    child: Option<Child>,
141}
142
143static MANAGED_PROCESSES: LazyLock<Mutex<HashMap<LocalProvider, ManagedProcess>>> =
144    LazyLock::new(|| Mutex::new(HashMap::new()));
145
146fn take_managed_child(provider: LocalProvider) -> Option<Child> {
147    let mut guard = MANAGED_PROCESSES.lock().unwrap_or_else(|e| e.into_inner());
148    guard.get_mut(&provider)?.child.take()
149}
150
151fn store_managed_child(provider: LocalProvider, child: Child) {
152    if let Ok(mut guard) = MANAGED_PROCESSES.lock() {
153        guard.entry(provider).or_insert_with(|| ManagedProcess { child: None }).child = Some(child);
154    }
155}
156
157fn is_managed_running(provider: LocalProvider) -> bool {
158    MANAGED_PROCESSES
159        .lock()
160        .ok()
161        .and_then(|guard| guard.get(&provider).map(|p| p.child.is_some()))
162        .unwrap_or(false)
163}
164
165// ---------------------------------------------------------------------------
166// Public API
167// ---------------------------------------------------------------------------
168
169pub async fn probe_all() -> Vec<LocalServerStatus> {
170    let mut statuses = Vec::with_capacity(LocalProvider::all().len());
171    for &provider in LocalProvider::all() {
172        statuses.push(probe(provider).await);
173    }
174    statuses
175}
176
177pub async fn probe(provider: LocalProvider) -> LocalServerStatus {
178    match provider {
179        LocalProvider::Ollama => probe_ollama().await,
180        LocalProvider::LmStudio => probe_lmstudio().await,
181        LocalProvider::LlamaCpp => probe_llamacpp().await,
182    }
183}
184
185pub async fn start(provider: LocalProvider) -> Result<String> {
186    match provider {
187        LocalProvider::Ollama => start_ollama().await,
188        LocalProvider::LmStudio => start_lmstudio().await,
189        LocalProvider::LlamaCpp => start_llamacpp().await,
190    }
191}
192
193pub async fn stop(provider: LocalProvider) -> Result<String> {
194    match provider {
195        LocalProvider::Ollama => stop_ollama().await,
196        LocalProvider::LmStudio => stop_lmstudio().await,
197        LocalProvider::LlamaCpp => stop_llamacpp().await,
198    }
199}
200
201pub fn capabilities(provider: LocalProvider) -> LocalServerCapabilities {
202    match provider {
203        LocalProvider::Ollama => caps_ollama(),
204        LocalProvider::LmStudio => caps_lmstudio(),
205        LocalProvider::LlamaCpp => caps_llamacpp(),
206    }
207}
208
209pub fn env_config(provider: LocalProvider) -> Vec<EnvVarInfo> {
210    match provider {
211        LocalProvider::Ollama => vec![EnvVarInfo {
212            name: env_vars::OLLAMA_BASE_URL,
213            current_value: std::env::var(env_vars::OLLAMA_BASE_URL).ok(),
214            description: "Ollama server base URL (default: http://localhost:11434)",
215        }],
216        LocalProvider::LmStudio => vec![EnvVarInfo {
217            name: env_vars::LMSTUDIO_BASE_URL,
218            current_value: std::env::var(env_vars::LMSTUDIO_BASE_URL).ok(),
219            description: "LM Studio server base URL (default: http://localhost:1234/v1)",
220        }],
221        LocalProvider::LlamaCpp => vec![
222            EnvVarInfo {
223                name: env_vars::LLAMACPP_BASE_URL,
224                current_value: std::env::var(env_vars::LLAMACPP_BASE_URL).ok(),
225                description: "llama.cpp server base URL (default: http://localhost:8080/v1)",
226            },
227            EnvVarInfo {
228                name: env_vars::LLAMACPP_MODEL_PATH,
229                current_value: std::env::var(env_vars::LLAMACPP_MODEL_PATH).ok(),
230                description: "Path to .gguf model file for auto-start",
231            },
232            EnvVarInfo {
233                name: env_vars::LLAMACPP_BINARY_PATH,
234                current_value: std::env::var(env_vars::LLAMACPP_BINARY_PATH).ok(),
235                description: "Path to llama-server binary (default: search PATH)",
236            },
237            EnvVarInfo {
238                name: env_vars::LLAMACPP_EXTRA_ARGS,
239                current_value: std::env::var(env_vars::LLAMACPP_EXTRA_ARGS).ok(),
240                description: "Extra arguments passed to llama-server",
241            },
242        ],
243    }
244}
245
246pub fn troubleshoot(status: &LocalServerStatus, caps: &LocalServerCapabilities) -> Vec<String> {
247    let mut lines = Vec::new();
248    lines.push(format!("{} Troubleshoot", status.provider.display_name()));
249    lines.push(String::new());
250
251    if status.running {
252        lines.push("Server is running and responding.".to_string());
253        if status.available_models.is_empty() {
254            lines.push("No models are currently available.".to_string());
255            match status.provider {
256                LocalProvider::Ollama => {
257                    lines.push("  Pull a model: ollama pull gemma3".to_string());
258                }
259                LocalProvider::LmStudio => {
260                    lines.push("  Download a model in LM Studio or run: lms get <model>".to_string());
261                }
262                LocalProvider::LlamaCpp => {
263                    lines.push(format!("  Set {}=/path/to/model.gguf and restart", env_vars::LLAMACPP_MODEL_PATH));
264                }
265            }
266        }
267        return lines;
268    }
269
270    lines.push("Status: Not running".to_string());
271    if let Some(err) = &status.error {
272        lines.push(format!("Error: {err}"));
273    }
274    lines.push(String::new());
275
276    match status.provider {
277        LocalProvider::Ollama => {
278            if !caps.binary_found {
279                lines.push("Ollama is not installed.".to_string());
280                lines.push("  Install: brew install ollama".to_string());
281                lines.push("  Or: https://github.com/ollama/ollama?tab=readme-ov-file".to_string());
282            } else {
283                lines.push("Ollama is installed but the server is not running.".to_string());
284                lines.push("  Start: ollama serve".to_string());
285                lines.push("  Or: /local start ollama".to_string());
286            }
287            lines.push("  Logs: ~/.ollama/logs/server.log".to_string());
288        }
289        LocalProvider::LmStudio => {
290            if !caps.binary_found {
291                lines.push("LM Studio CLI (lms) not found.".to_string());
292                lines.push("  Install LM Studio: https://lmstudio.ai/download".to_string());
293                lines.push("  The lms CLI ships with LM Studio.".to_string());
294            } else {
295                lines.push("LM Studio server is not running.".to_string());
296                lines.push("  Start: lms server start".to_string());
297                lines.push("  Or: /local start lmstudio".to_string());
298                lines.push("  Status: lms server status --json".to_string());
299            }
300        }
301        LocalProvider::LlamaCpp => {
302            if !caps.binary_found {
303                lines.push("llama-server binary not found.".to_string());
304                lines.push("  Install: https://llama.app".to_string());
305                lines.push(format!("  Or set {}=/path/to/llama-server", env_vars::LLAMACPP_BINARY_PATH));
306            } else {
307                lines.push("llama.cpp server is not running.".to_string());
308                let model_path = std::env::var(env_vars::LLAMACPP_MODEL_PATH).ok();
309                if model_path.is_none() {
310                    lines.push(format!("  Set {}=/path/to/model.gguf for auto-start", env_vars::LLAMACPP_MODEL_PATH));
311                }
312                lines.push("  Or: /local start llamacpp".to_string());
313            }
314        }
315    }
316
317    lines
318}
319
320// ---------------------------------------------------------------------------
321// Probe implementations
322// ---------------------------------------------------------------------------
323
324async fn probe_ollama() -> LocalServerStatus {
325    let base = LocalProvider::Ollama.host_root();
326    let client = vtcode_commons::http::create_client_with_timeout(PROBE_TIMEOUT);
327
328    // Check /api/tags for availability + models
329    let tags_url = format!("{}/api/tags", base.trim_end_matches('/'));
330    let tags_resp = match client.get(&tags_url).send().await {
331        Ok(resp) => resp,
332        Err(e) => {
333            let mut s = LocalServerStatus::not_running(LocalProvider::Ollama, e.to_string());
334            if is_managed_running(LocalProvider::Ollama) {
335                s.error = Some("Managed process exists but server not responding yet".into());
336            }
337            return s;
338        }
339    };
340
341    if !tags_resp.status().is_success() {
342        return LocalServerStatus::not_running(LocalProvider::Ollama, format!("HTTP {}", tags_resp.status()));
343    }
344
345    let available_models = tags_resp
346        .json::<OllamaTagsResponse>()
347        .await
348        .map(|r| r.models.into_iter().map(|m| m.name).collect())
349        .unwrap_or_default();
350
351    // Check /api/ps for running models
352    let ps_url = format!("{}/api/ps", base.trim_end_matches('/'));
353    let running_models = parse_json_opt::<OllamaPsResponse>(client.get(&ps_url).send().await.ok())
354        .await
355        .map(|r| r.models.into_iter().map(|m| m.name).collect())
356        .unwrap_or_default();
357
358    // Check /api/version
359    let version_url = format!("{}/api/version", base.trim_end_matches('/'));
360    let version = parse_json_opt::<OllamaVersionResponse>(client.get(&version_url).send().await.ok())
361        .await
362        .and_then(|r| r.version);
363
364    LocalServerStatus {
365        provider: LocalProvider::Ollama,
366        running: true,
367        endpoint: LocalProvider::Ollama.base_url(),
368        available_models,
369        running_models,
370        version,
371        error: None,
372    }
373}
374
375async fn probe_lmstudio() -> LocalServerStatus {
376    let base = LocalProvider::LmStudio.base_url();
377    let client = vtcode_commons::http::create_client_with_timeout(PROBE_TIMEOUT);
378
379    let models_url = format!("{}/models", base.trim_end_matches('/'));
380    let resp = match client.get(&models_url).send().await {
381        Ok(resp) => resp,
382        Err(e) => return LocalServerStatus::not_running(LocalProvider::LmStudio, e.to_string()),
383    };
384
385    if !resp.status().is_success() {
386        return LocalServerStatus::not_running(LocalProvider::LmStudio, format!("HTTP {}", resp.status()));
387    }
388
389    let available_models = resp
390        .json::<LmStudioModelsResponse>()
391        .await
392        .map(|r| r.data.into_iter().map(|m| m.id).collect())
393        .unwrap_or_default();
394
395    LocalServerStatus {
396        provider: LocalProvider::LmStudio,
397        running: true,
398        endpoint: LocalProvider::LmStudio.base_url(),
399        available_models,
400        running_models: Vec::new(),
401        version: None,
402        error: None,
403    }
404}
405
406async fn probe_llamacpp() -> LocalServerStatus {
407    let base = LocalProvider::LlamaCpp.host_root();
408    let client = vtcode_commons::http::create_client_with_timeout(PROBE_TIMEOUT);
409
410    // Check /health
411    let health_url = format!("{}/health", base.trim_end_matches('/'));
412    let health_resp = match client.get(&health_url).send().await {
413        Ok(resp) => resp,
414        Err(e) => return LocalServerStatus::not_running(LocalProvider::LlamaCpp, e.to_string()),
415    };
416
417    if !health_resp.status().is_success() {
418        return LocalServerStatus::not_running(LocalProvider::LlamaCpp, format!("HTTP {}", health_resp.status()));
419    }
420
421    // Check /models
422    let models_url = format!("{}/models", base.trim_end_matches('/'));
423    let available_models = parse_json_opt::<LlamaCppModelsResponse>(client.get(&models_url).send().await.ok())
424        .await
425        .map(|r| r.data.into_iter().map(|m| m.id).collect())
426        .unwrap_or_default();
427
428    LocalServerStatus {
429        provider: LocalProvider::LlamaCpp,
430        running: true,
431        endpoint: LocalProvider::LlamaCpp.base_url(),
432        available_models,
433        running_models: Vec::new(),
434        version: None,
435        error: None,
436    }
437}
438
439// ---------------------------------------------------------------------------
440// Start implementations
441// ---------------------------------------------------------------------------
442
443async fn start_ollama() -> Result<String> {
444    let caps = caps_ollama();
445    if !caps.binary_found {
446        anyhow::bail!(
447            "Ollama is not installed. Install with: brew install ollama\n\
448             Or visit: https://github.com/ollama/ollama?tab=readme-ov-file"
449        );
450    }
451
452    // Check if already running
453    let status = probe_ollama().await;
454    if status.running {
455        return Ok("Ollama is already running.".to_string());
456    }
457
458    let binary = caps.binary_path.unwrap_or_else(|| "ollama".to_string());
459    let mut cmd = tokio::process::Command::new(&binary);
460    sanitize_tokio_command_environment(&mut cmd, &[]);
461    cmd.arg("serve")
462        .stdin(Stdio::null())
463        .stdout(Stdio::null())
464        .stderr(Stdio::null())
465        .kill_on_drop(true);
466
467    let child = cmd
468        .spawn()
469        .with_context(|| format!("Failed to start Ollama with `{binary} serve`"))?;
470
471    store_managed_child(LocalProvider::Ollama, child);
472
473    // Wait for it to become ready
474    wait_for_ready(LocalProvider::Ollama, Duration::from_secs(10)).await?;
475
476    Ok("Ollama server started.".to_string())
477}
478
479async fn start_lmstudio() -> Result<String> {
480    let caps = caps_lmstudio();
481    if !caps.binary_found {
482        anyhow::bail!(
483            "LM Studio CLI (lms) not found.\n\
484             Install LM Studio from https://lmstudio.ai/download\n\
485             The lms CLI ships with the app."
486        );
487    }
488
489    // Check if already running
490    let status = probe_lmstudio().await;
491    if status.running {
492        return Ok("LM Studio server is already running.".to_string());
493    }
494
495    let binary = caps.binary_path.unwrap_or_else(|| "lms".to_string());
496    let mut command = tokio::process::Command::new(&binary);
497    sanitize_tokio_command_environment(&mut command, &[]);
498    let output = command
499        .args(["server", "start"])
500        .output()
501        .await
502        .with_context(|| format!("Failed to run `{binary} server start`"))?;
503
504    if !output.status.success() {
505        let stderr = String::from_utf8_lossy(&output.stderr);
506        anyhow::bail!("lms server start failed: {}", stderr.trim());
507    }
508
509    // Wait for it to become ready
510    wait_for_ready(LocalProvider::LmStudio, Duration::from_secs(10)).await?;
511
512    Ok("LM Studio server started.".to_string())
513}
514
515async fn start_llamacpp() -> Result<String> {
516    let caps = caps_llamacpp();
517    if !caps.binary_found {
518        anyhow::bail!(
519            "llama-server binary not found.\n\
520             Install from https://llama.app\n\
521             Or set {}=/path/to/llama-server",
522            env_vars::LLAMACPP_BINARY_PATH
523        );
524    }
525
526    // Check if already running
527    let status = probe_llamacpp().await;
528    if status.running {
529        return Ok("llama.cpp server is already running.".to_string());
530    }
531
532    let model_path = std::env::var(env_vars::LLAMACPP_MODEL_PATH)
533        .ok()
534        .filter(|v| !v.trim().is_empty());
535    let model_path = match model_path {
536        Some(path) => path,
537        None => anyhow::bail!(
538            "Set {}=/path/to/model.gguf to enable auto-start for llama.cpp",
539            env_vars::LLAMACPP_MODEL_PATH
540        ),
541    };
542
543    let binary = caps.binary_path.unwrap_or_else(|| "llama-server".to_string());
544    let port = extract_port(&LocalProvider::LlamaCpp.base_url()).unwrap_or(LocalProvider::LlamaCpp.default_port());
545
546    let mut args = vec!["-m".to_string(), model_path, "--port".to_string(), port.to_string()];
547    if let Ok(extra) = std::env::var(env_vars::LLAMACPP_EXTRA_ARGS)
548        && !extra.trim().is_empty()
549    {
550        args.extend(shell_words::split(&extra).unwrap_or_default());
551    }
552
553    let mut cmd = tokio::process::Command::new(&binary);
554    sanitize_tokio_command_environment(&mut cmd, &[]);
555    cmd.args(&args)
556        .stdin(Stdio::null())
557        .stdout(Stdio::null())
558        .stderr(Stdio::null())
559        .kill_on_drop(true);
560
561    let child = cmd
562        .spawn()
563        .with_context(|| format!("Failed to start llama-server (`{binary} -m <model> --port {port}`)"))?;
564
565    store_managed_child(LocalProvider::LlamaCpp, child);
566
567    // Wait for it to become ready
568    wait_for_ready(LocalProvider::LlamaCpp, Duration::from_secs(30)).await?;
569
570    Ok("llama.cpp server started.".to_string())
571}
572
573// ---------------------------------------------------------------------------
574// Stop implementations
575// ---------------------------------------------------------------------------
576
577async fn stop_ollama() -> Result<String> {
578    if let Some(mut child) = take_managed_child(LocalProvider::Ollama) {
579        child.kill().await.ok();
580        return Ok("Ollama server stopped.".to_string());
581    }
582
583    // No managed process; check if it's running externally
584    let status = probe_ollama().await;
585    if !status.running {
586        return Ok("Ollama is not running.".to_string());
587    }
588
589    anyhow::bail!(
590        "Ollama is running but was not started by VT Code.\n\
591         Stop it manually or kill the process."
592    )
593}
594
595async fn stop_lmstudio() -> Result<String> {
596    let caps = caps_lmstudio();
597    if !caps.binary_found {
598        return Ok("LM Studio CLI not found; nothing to stop.".to_string());
599    }
600
601    let status = probe_lmstudio().await;
602    if !status.running {
603        return Ok("LM Studio server is not running.".to_string());
604    }
605
606    let binary = caps.binary_path.unwrap_or_else(|| "lms".to_string());
607    let mut command = tokio::process::Command::new(&binary);
608    sanitize_tokio_command_environment(&mut command, &[]);
609    let output = command
610        .args(["server", "stop"])
611        .output()
612        .await
613        .with_context(|| format!("Failed to run `{binary} server stop`"))?;
614
615    if !output.status.success() {
616        let stderr = String::from_utf8_lossy(&output.stderr);
617        anyhow::bail!("lms server stop failed: {}", stderr.trim());
618    }
619
620    Ok("LM Studio server stopped.".to_string())
621}
622
623async fn stop_llamacpp() -> Result<String> {
624    // Try our own managed child first
625    if let Some(mut child) = take_managed_child(LocalProvider::LlamaCpp) {
626        child.kill().await.ok();
627        return Ok("llama.cpp server stopped.".to_string());
628    }
629
630    // Check if running externally
631    let status = probe_llamacpp().await;
632    if !status.running {
633        return Ok("llama.cpp server is not running.".to_string());
634    }
635
636    anyhow::bail!(
637        "llama.cpp is running but was not started by VT Code.\n\
638         Stop it manually or kill the process."
639    )
640}
641
642// ---------------------------------------------------------------------------
643// Capabilities
644// ---------------------------------------------------------------------------
645
646fn caps_ollama() -> LocalServerCapabilities {
647    let (found, path) = find_binary("ollama");
648    LocalServerCapabilities {
649        can_start: found,
650        can_stop: is_managed_running(LocalProvider::Ollama),
651        binary_found: found,
652        binary_name: "ollama",
653        binary_path: path,
654    }
655}
656
657fn caps_lmstudio() -> LocalServerCapabilities {
658    // Try `lms` on PATH, then fallback to ~/.lmstudio/bin/lms
659    let (found, path) = find_binary("lms");
660    let (found, path) = if !found {
661        find_lms_fallback().unwrap_or((false, None))
662    } else {
663        (found, path)
664    };
665    LocalServerCapabilities {
666        can_start: found,
667        can_stop: found,
668        binary_found: found,
669        binary_name: "lms",
670        binary_path: path,
671    }
672}
673
674fn caps_llamacpp() -> LocalServerCapabilities {
675    // Check LLAMACPP_BINARY_PATH first, then PATH
676    let (found, path) = if let Ok(explicit) = std::env::var(env_vars::LLAMACPP_BINARY_PATH) {
677        if !explicit.trim().is_empty() && std::path::Path::new(explicit.trim()).exists() {
678            (true, Some(explicit.trim().to_string()))
679        } else {
680            find_binary("llama-server")
681        }
682    } else {
683        find_binary("llama-server")
684    };
685
686    LocalServerCapabilities {
687        can_start: found
688            && std::env::var(env_vars::LLAMACPP_MODEL_PATH)
689                .ok()
690                .filter(|v| !v.trim().is_empty())
691                .is_some(),
692        can_stop: is_managed_running(LocalProvider::LlamaCpp),
693        binary_found: found,
694        binary_name: "llama-server",
695        binary_path: path,
696    }
697}
698
699// ---------------------------------------------------------------------------
700// Helpers
701// ---------------------------------------------------------------------------
702
703/// Returns true when `base_url` points at a loopback/local endpoint.
704///
705/// Shared by local providers so the "is this localhost?" check lives in one
706/// place. Accepts `localhost`, `127.0.0.1` (and the whole `127.0.0.0/8`
707/// subnet via prefix match), `::1`, and `0.0.0.0` over http/https.
708pub(crate) fn is_local_base_url(base_url: &str) -> bool {
709    let lowered = base_url.trim().to_ascii_lowercase();
710    const LOCAL_PREFIXES: &[&str] = &[
711        "http://localhost",
712        "https://localhost",
713        "http://127.",
714        "https://127.",
715        "http://0.0.0.0",
716        "https://0.0.0.0",
717        "http://[::1]",
718        "https://[::1]",
719    ];
720    if LOCAL_PREFIXES.iter().any(|prefix| lowered.starts_with(*prefix)) {
721        return true;
722    }
723    if let Ok(parsed) = url::Url::parse(lowered.trim_end_matches('/'))
724        && let Some(host) = parsed.host_str()
725    {
726        return matches!(host, "localhost" | "127.0.0.1" | "::1" | "0.0.0.0");
727    }
728    false
729}
730
731fn strip_path_suffix(url: &str) -> String {
732    // Strip /v1 or /api/v1 suffix to get the host root
733    let trimmed = url.trim_end_matches('/');
734    if let Some(pos) = trimmed.rfind("/v1") {
735        trimmed[..pos].to_string()
736    } else {
737        trimmed.to_string()
738    }
739}
740
741fn extract_port(url: &str) -> Option<u16> {
742    let stripped = strip_path_suffix(url);
743    url::Url::parse(&stripped).ok().and_then(|u| u.port())
744}
745
746fn find_binary(name: &str) -> (bool, Option<String>) {
747    which::which(name)
748        .map(|p| (true, Some(p.to_string_lossy().into_owned())))
749        .unwrap_or((false, None))
750}
751
752fn find_lms_fallback() -> Option<(bool, Option<String>)> {
753    let home = std::env::var("HOME").ok()?;
754    let fallback = format!("{home}/.lmstudio/bin/lms");
755    if std::path::Path::new(&fallback).exists() {
756        Some((true, Some(fallback)))
757    } else {
758        None
759    }
760}
761
762async fn wait_for_ready(provider: LocalProvider, timeout: Duration) -> Result<()> {
763    let deadline = tokio::time::Instant::now() + timeout;
764    let mut last_error = String::new();
765
766    while tokio::time::Instant::now() < deadline {
767        let status = probe(provider).await;
768        if status.running {
769            return Ok(());
770        }
771        if let Some(err) = &status.error {
772            last_error = err.clone();
773        }
774        tokio::time::sleep(Duration::from_millis(500)).await;
775    }
776
777    anyhow::bail!(
778        "Timed out waiting for {} to start after {}s. Last: {}",
779        provider.display_name(),
780        timeout.as_secs(),
781        last_error
782    )
783}
784
785// Response types
786
787#[derive(Deserialize)]
788struct OllamaTagsResponse {
789    models: Vec<OllamaModelSummary>,
790}
791
792#[derive(Deserialize)]
793struct OllamaModelSummary {
794    name: String,
795}
796
797#[derive(Deserialize)]
798struct OllamaPsResponse {
799    models: Vec<OllamaRunningModel>,
800}
801
802#[derive(Deserialize)]
803struct OllamaRunningModel {
804    name: String,
805}
806
807#[derive(Deserialize)]
808struct OllamaVersionResponse {
809    version: Option<String>,
810}
811
812#[derive(Deserialize)]
813struct LmStudioModelsResponse {
814    data: Vec<LmStudioModel>,
815}
816
817#[derive(Deserialize)]
818struct LmStudioModel {
819    id: String,
820}
821
822#[derive(Deserialize)]
823struct LlamaCppModelsResponse {
824    data: Vec<LlamaCppModel>,
825}
826
827#[derive(Deserialize)]
828struct LlamaCppModel {
829    id: String,
830}
831
832// Helper: parse response as JSON, returning None on failure
833async fn parse_json_opt<T: serde::de::DeserializeOwned>(resp: Option<reqwest::Response>) -> Option<T> {
834    let resp = resp?;
835    if !resp.status().is_success() {
836        return None;
837    }
838    resp.json::<T>().await.ok()
839}
840
841#[cfg(test)]
842mod tests {
843    use super::*;
844
845    #[test]
846    fn test_provider_from_key() {
847        assert_eq!(LocalProvider::from_key("ollama"), Some(LocalProvider::Ollama));
848        assert_eq!(LocalProvider::from_key("lmstudio"), Some(LocalProvider::LmStudio));
849        assert_eq!(LocalProvider::from_key("lm-studio"), Some(LocalProvider::LmStudio));
850        assert_eq!(LocalProvider::from_key("llamacpp"), Some(LocalProvider::LlamaCpp));
851        assert_eq!(LocalProvider::from_key("llama.cpp"), Some(LocalProvider::LlamaCpp));
852        assert_eq!(LocalProvider::from_key("unknown"), None);
853    }
854
855    #[test]
856    fn test_provider_key_roundtrip() {
857        for &p in LocalProvider::all() {
858            assert_eq!(LocalProvider::from_key(p.key()), Some(p));
859        }
860    }
861
862    #[test]
863    fn test_provider_display_names() {
864        assert_eq!(LocalProvider::Ollama.display_name(), "Ollama");
865        assert_eq!(LocalProvider::LmStudio.display_name(), "LM Studio");
866        assert_eq!(LocalProvider::LlamaCpp.display_name(), "llama.cpp");
867    }
868
869    #[test]
870    fn test_strip_path_suffix() {
871        assert_eq!(strip_path_suffix("http://localhost:11434/v1"), "http://localhost:11434");
872        assert_eq!(strip_path_suffix("http://localhost:1234/v1/"), "http://localhost:1234");
873        assert_eq!(strip_path_suffix("http://localhost:8080"), "http://localhost:8080");
874    }
875
876    #[test]
877    fn test_is_local_base_url_accepts_loopback() {
878        for url in [
879            "http://localhost:11434",
880            "https://localhost:1234/v1",
881            "http://127.0.0.1:8080/v1",
882            "http://127.1.2.3:9999",
883            "http://0.0.0.0:8080",
884            "http://[::1]:1234/v1",
885        ] {
886            assert!(is_local_base_url(url), "expected local: {url}");
887        }
888    }
889
890    #[test]
891    fn test_is_local_base_url_rejects_remote() {
892        for url in [
893            "http://192.168.1.10:11434",
894            "https://api.openai.com/v1",
895            "http://example.com:8080/v1",
896            "http://10.0.0.5:1234",
897        ] {
898            assert!(!is_local_base_url(url), "expected remote: {url}");
899        }
900    }
901}
vtcode_llm/providers/local_server.rs

vtcode_llm/providers/
local_server.rs
