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

vtcode_core/llm/providers/
local_server.rs
