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

vtcode_llm/providers/
local_server.rs
