newt_core/
dgx.rs

1//! NVIDIA DGX endpoint configuration — the data model behind the
2//! `newt dgx` command suite.
3//!
4//! A [`DgxConfig`] describes one or more DGX *nodes*, each exposing up to four
5//! named *endpoints*, a set of saved *formations* (model + endpoint presets),
6//! and the currently-active selection. It is an optional sub-table of
7//! [`crate::Config`] (`[dgx]` in `newt.toml`); when absent, newt never contacts
8//! a DGX host unless a `NEWT_DGX_*` environment variable is set — there are no
9//! leaky defaults.
10//!
11//! # Endpoint flavors ([`EndpointKind`])
12//!
13//! | kind | config key | example URL |
14//! |------|-----------|-------------|
15//! | direct DGX (Traefik) | `ollama` | `https://REDACTED-HOST` |
16//! | round-robin LB | `ollama_lb` | `https://REDACTED-HOST` |
17//! | in-cluster proxy | `in_cluster` | `http://ollama-proxy.inference.svc.cluster.local:11434` |
18//! | vLLM (OpenAI-compatible) | `vllm` | `http://REDACTED-HOST:8000` |
19//!
20//! The first three speak Ollama's native `/api` surface; only `vllm` is
21//! OpenAI-compatible (`/v1`). HTTPS endpoints (the Traefik-fronted hosts on
22//! port 443) are supported directly — give the full `https://...` URL with no
23//! port.
24//!
25//! # Environment overrides
26//!
27//! Resolution consults the environment before the config file, so an
28//! unconfigured install can still target a DGX:
29//!
30//! - `NEWT_DGX_OLLAMA_URL`, `NEWT_DGX_OLLAMA_LB_URL`,
31//!   `NEWT_DGX_IN_CLUSTER_URL`, `NEWT_DGX_VLLM_URL` — full URL for one flavor.
32//! - `NEWT_DGX_HOST` (plus optional `NEWT_DGX_SCHEME`, `NEWT_DGX_OLLAMA_PORT`,
33//!   `NEWT_DGX_VLLM_PORT`) — synthesize the `ollama`/`vllm` URL from a bare
34//!   host when no explicit URL is set. The LB and in-cluster proxy are distinct
35//!   hostnames, not ports on `NEWT_DGX_HOST`, so they are never synthesized.
36//! - `NEWT_DGX_MODEL` — active model id.
37//! - `NEWT_DGX_SSH_HOST`, `NEWT_DGX_SSH_USER` — SSH target for `run`/`push`.
38
39use std::fmt;
40use std::str::FromStr;
41
42use serde::{Deserialize, Serialize};
43use thiserror::Error;
44
45// ---------------------------------------------------------------------------
46// Endpoint kind
47// ---------------------------------------------------------------------------
48
49/// The four endpoint flavors a DGX node can expose.
50#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default, Serialize, Deserialize)]
51#[serde(rename_all = "snake_case")]
52pub enum EndpointKind {
53    /// Direct DGX Ollama endpoint (typically Traefik-fronted HTTPS).
54    #[default]
55    Ollama,
56    /// Round-robin load balancer across Ollama backends.
57    OllamaLb,
58    /// In-cluster, model-aware Ollama proxy (for pods).
59    InCluster,
60    /// vLLM OpenAI-compatible endpoint.
61    Vllm,
62}
63
64impl EndpointKind {
65    /// Every endpoint kind, in display order.
66    pub const ALL: [Self; 4] = [Self::Ollama, Self::OllamaLb, Self::InCluster, Self::Vllm];
67
68    /// The TOML config key / CLI token for this kind.
69    pub fn as_str(self) -> &'static str {
70        match self {
71            Self::Ollama => "ollama",
72            Self::OllamaLb => "ollama_lb",
73            Self::InCluster => "in_cluster",
74            Self::Vllm => "vllm",
75        }
76    }
77
78    /// The environment variable that overrides this flavor's full URL.
79    pub fn url_env_var(self) -> &'static str {
80        match self {
81            Self::Ollama => "NEWT_DGX_OLLAMA_URL",
82            Self::OllamaLb => "NEWT_DGX_OLLAMA_LB_URL",
83            Self::InCluster => "NEWT_DGX_IN_CLUSTER_URL",
84            Self::Vllm => "NEWT_DGX_VLLM_URL",
85        }
86    }
87
88    /// Whether this flavor speaks the OpenAI-compatible (`/v1`) API. Only
89    /// `vllm` does; the rest speak Ollama's native `/api`.
90    pub fn is_openai_compatible(self) -> bool {
91        matches!(self, Self::Vllm)
92    }
93}
94
95impl fmt::Display for EndpointKind {
96    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
97        f.write_str(self.as_str())
98    }
99}
100
101impl FromStr for EndpointKind {
102    type Err = String;
103
104    fn from_str(s: &str) -> Result<Self, Self::Err> {
105        let norm = s.trim().to_ascii_lowercase().replace('-', "_");
106        match norm.as_str() {
107            "ollama" => Ok(Self::Ollama),
108            "ollama_lb" | "lb" => Ok(Self::OllamaLb),
109            "in_cluster" | "incluster" | "cluster" | "proxy" => Ok(Self::InCluster),
110            "vllm" => Ok(Self::Vllm),
111            other => Err(format!(
112                "unknown endpoint kind {other:?} (expected one of: ollama, ollama_lb, in_cluster, vllm)"
113            )),
114        }
115    }
116}
117
118// ---------------------------------------------------------------------------
119// Errors
120// ---------------------------------------------------------------------------
121
122/// Why a DGX endpoint / model / host could not be resolved. Every variant is
123/// actionable: it names what is missing and how to set it.
124#[derive(Debug, Clone, PartialEq, Eq, Error)]
125pub enum DgxNotConfigured {
126    /// No nodes in config and no `NEWT_DGX_*` fallback.
127    #[error("no DGX nodes configured — set NEWT_DGX_HOST=<host>  or  NEWT_DGX_OLLAMA_URL=<url>")]
128    NoNodes,
129
130    /// Several nodes exist but none is marked active.
131    #[error(
132        "no active DGX node selected ({count} configured) — set [dgx].active_node in your config"
133    )]
134    NoActiveNode { count: usize },
135
136    /// `active_node` names a node that isn't in the list.
137    #[error("DGX node {name:?} not found in config")]
138    NodeNotFound { name: String },
139
140    /// The active node has no URL for the requested flavor.
141    #[error("DGX node {node:?} has no {kind} endpoint set")]
142    EndpointUnset { node: String, kind: EndpointKind },
143
144    /// No active model and no `NEWT_DGX_MODEL`.
145    #[error("no active DGX model — set [dgx].active_model or NEWT_DGX_MODEL=<model-id>")]
146    NoActiveModel,
147
148    /// No SSH host for `run`/`push`.
149    #[error("DGX node {node:?} has no ssh_host — set it in config or via NEWT_DGX_SSH_HOST")]
150    NoSshHost { node: String },
151}
152
153// ---------------------------------------------------------------------------
154// Node / formation
155// ---------------------------------------------------------------------------
156
157/// A single DGX host and the endpoints it exposes.
158#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
159pub struct DgxNode {
160    /// Short name used by `newt dgx node use <name>`. For a per-file drop-in
161    /// (`~/.newt/dgx/<name>.toml`) this is overwritten by the filename stem, so
162    /// the file body may omit it.
163    #[serde(default)]
164    pub name: String,
165
166    /// Direct DGX Ollama URL (e.g. `https://REDACTED-HOST`).
167    #[serde(default, skip_serializing_if = "Option::is_none")]
168    pub ollama: Option<String>,
169
170    /// Round-robin LB URL (e.g. `https://REDACTED-HOST`).
171    #[serde(default, skip_serializing_if = "Option::is_none")]
172    pub ollama_lb: Option<String>,
173
174    /// In-cluster proxy URL (plain HTTP cluster DNS).
175    #[serde(default, skip_serializing_if = "Option::is_none")]
176    pub in_cluster: Option<String>,
177
178    /// vLLM OpenAI-compatible URL.
179    #[serde(default, skip_serializing_if = "Option::is_none")]
180    pub vllm: Option<String>,
181
182    /// SSH host for `run`/`push` (may differ from the HTTP hostnames).
183    #[serde(default, skip_serializing_if = "Option::is_none")]
184    pub ssh_host: Option<String>,
185
186    /// SSH user; falls back to `$USER`, then `dgx`.
187    #[serde(default, skip_serializing_if = "Option::is_none")]
188    pub ssh_user: Option<String>,
189}
190
191impl DgxNode {
192    /// The configured URL for `kind` on this node, if any.
193    pub fn endpoint(&self, kind: EndpointKind) -> Option<&str> {
194        match kind {
195            EndpointKind::Ollama => self.ollama.as_deref(),
196            EndpointKind::OllamaLb => self.ollama_lb.as_deref(),
197            EndpointKind::InCluster => self.in_cluster.as_deref(),
198            EndpointKind::Vllm => self.vllm.as_deref(),
199        }
200    }
201}
202
203/// A named (model, endpoint) preset — `newt dgx formation <name>`.
204#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
205pub struct DgxFormation {
206    pub name: String,
207    pub model: String,
208    #[serde(default)]
209    pub endpoint: EndpointKind,
210}
211
212// ---------------------------------------------------------------------------
213// Top-level config
214// ---------------------------------------------------------------------------
215
216/// The `[dgx]` sub-table: nodes, formations, and the active selection.
217#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
218#[serde(default)]
219pub struct DgxConfig {
220    /// Name of the active node; `None` means "the only node, if unambiguous".
221    #[serde(skip_serializing_if = "Option::is_none")]
222    pub active_node: Option<String>,
223
224    /// Active endpoint flavor (defaults to `ollama`).
225    pub active_endpoint: EndpointKind,
226
227    /// Active model id.
228    #[serde(skip_serializing_if = "Option::is_none")]
229    pub active_model: Option<String>,
230
231    /// Configured DGX nodes.
232    #[serde(skip_serializing_if = "Vec::is_empty")]
233    pub nodes: Vec<DgxNode>,
234
235    /// Saved (model, endpoint) presets.
236    #[serde(skip_serializing_if = "Vec::is_empty")]
237    pub formations: Vec<DgxFormation>,
238}
239
240impl DgxConfig {
241    /// Look up a node by name.
242    pub fn node(&self, name: &str) -> Option<&DgxNode> {
243        self.nodes.iter().find(|n| n.name == name)
244    }
245
246    /// Look up a formation by name.
247    pub fn formation(&self, name: &str) -> Option<&DgxFormation> {
248        self.formations.iter().find(|f| f.name == name)
249    }
250
251    /// Resolve the active node: the one named by `active_node`, or — if that is
252    /// unset — the single configured node when there is exactly one.
253    pub fn active_node(&self) -> Result<&DgxNode, DgxNotConfigured> {
254        match &self.active_node {
255            Some(name) => self
256                .node(name)
257                .ok_or_else(|| DgxNotConfigured::NodeNotFound { name: name.clone() }),
258            None => match self.nodes.as_slice() {
259                [] => Err(DgxNotConfigured::NoNodes),
260                [only] => Ok(only),
261                many => Err(DgxNotConfigured::NoActiveNode { count: many.len() }),
262            },
263        }
264    }
265
266    /// Resolve the URL for the *active* endpoint flavor.
267    pub fn resolve_endpoint(&self) -> Result<String, DgxNotConfigured> {
268        self.resolve_endpoint_for(self.active_endpoint)
269    }
270
271    /// Resolve the URL for a specific endpoint flavor, consulting the real
272    /// process environment.
273    pub fn resolve_endpoint_for(&self, kind: EndpointKind) -> Result<String, DgxNotConfigured> {
274        self.resolve_endpoint_for_with(kind, &|k| std::env::var(k).ok())
275    }
276
277    /// [`resolve_endpoint_for`](Self::resolve_endpoint_for) with an injectable
278    /// environment lookup (for testing). Precedence: per-flavor URL env var →
279    /// configured node URL → `NEWT_DGX_HOST` synthesis (ollama/vllm only).
280    pub fn resolve_endpoint_for_with(
281        &self,
282        kind: EndpointKind,
283        get_env: &dyn Fn(&str) -> Option<String>,
284    ) -> Result<String, DgxNotConfigured> {
285        if let Some(url) = nonempty(get_env(kind.url_env_var())) {
286            return Ok(url);
287        }
288        if let Ok(node) = self.active_node() {
289            if let Some(url) = node.endpoint(kind) {
290                return Ok(url.to_string());
291            }
292        }
293        if let Some(url) = synth_from_host(kind, get_env) {
294            return Ok(url);
295        }
296        match self.active_node() {
297            Ok(node) => Err(DgxNotConfigured::EndpointUnset {
298                node: node.name.clone(),
299                kind,
300            }),
301            Err(e) => Err(e),
302        }
303    }
304
305    /// Resolve the active model id (env `NEWT_DGX_MODEL` wins).
306    pub fn resolve_active_model(&self) -> Result<String, DgxNotConfigured> {
307        self.resolve_active_model_with(&|k| std::env::var(k).ok())
308    }
309
310    /// [`resolve_active_model`](Self::resolve_active_model) with injectable env.
311    pub fn resolve_active_model_with(
312        &self,
313        get_env: &dyn Fn(&str) -> Option<String>,
314    ) -> Result<String, DgxNotConfigured> {
315        if let Some(model) = nonempty(get_env("NEWT_DGX_MODEL")) {
316            return Ok(model);
317        }
318        self.active_model
319            .clone()
320            .ok_or(DgxNotConfigured::NoActiveModel)
321    }
322
323    /// Resolve the SSH host (env `NEWT_DGX_SSH_HOST` wins).
324    pub fn ssh_host(&self) -> Result<String, DgxNotConfigured> {
325        self.ssh_host_with(&|k| std::env::var(k).ok())
326    }
327
328    /// [`ssh_host`](Self::ssh_host) with injectable env.
329    pub fn ssh_host_with(
330        &self,
331        get_env: &dyn Fn(&str) -> Option<String>,
332    ) -> Result<String, DgxNotConfigured> {
333        if let Some(host) = nonempty(get_env("NEWT_DGX_SSH_HOST")) {
334            return Ok(host);
335        }
336        let node = self.active_node()?;
337        node.ssh_host
338            .clone()
339            .ok_or_else(|| DgxNotConfigured::NoSshHost {
340                node: node.name.clone(),
341            })
342    }
343
344    /// Resolve the SSH user: `NEWT_DGX_SSH_USER` → node `ssh_user` → `$USER`
345    /// → `dgx`. Never fails.
346    pub fn ssh_user(&self) -> String {
347        self.ssh_user_with(&|k| std::env::var(k).ok())
348    }
349
350    /// [`ssh_user`](Self::ssh_user) with injectable env.
351    pub fn ssh_user_with(&self, get_env: &dyn Fn(&str) -> Option<String>) -> String {
352        if let Some(user) = nonempty(get_env("NEWT_DGX_SSH_USER")) {
353            return user;
354        }
355        if let Ok(node) = self.active_node() {
356            if let Some(user) = &node.ssh_user {
357                return user.clone();
358            }
359        }
360        nonempty(get_env("USER")).unwrap_or_else(|| "dgx".to_string())
361    }
362
363    /// A ready-to-write template for the reference DGX topology.
364    /// `newt dgx setup` offers this as the suggested starting point; it is
365    /// never applied automatically.
366    pub fn home_template() -> Self {
367        Self {
368            active_node: Some("home".to_string()),
369            active_endpoint: EndpointKind::Ollama,
370            active_model: Some("qwen2.5-coder:32b".to_string()),
371            nodes: vec![DgxNode {
372                name: "home".to_string(),
373                ollama: Some("https://REDACTED-HOST".to_string()),
374                ollama_lb: Some("https://REDACTED-HOST".to_string()),
375                in_cluster: Some(
376                    "http://ollama-proxy.inference.svc.cluster.local:11434".to_string(),
377                ),
378                vllm: None,
379                ssh_host: Some("REDACTED-HOST".to_string()),
380                ssh_user: None,
381            }],
382            formations: vec![
383                DgxFormation {
384                    name: "coding".to_string(),
385                    model: "qwen2.5-coder:32b".to_string(),
386                    endpoint: EndpointKind::Ollama,
387                },
388                DgxFormation {
389                    name: "review".to_string(),
390                    model: "llama3.1:70b".to_string(),
391                    endpoint: EndpointKind::InCluster,
392                },
393            ],
394        }
395    }
396}
397
398// ---------------------------------------------------------------------------
399// Helpers
400// ---------------------------------------------------------------------------
401
402/// Treat empty / whitespace-only values as unset.
403fn nonempty(value: Option<String>) -> Option<String> {
404    value.filter(|s| !s.trim().is_empty())
405}
406
407/// Synthesize an `ollama`/`vllm` URL from `NEWT_DGX_HOST` + optional scheme and
408/// port. Returns `None` for `ollama_lb` / `in_cluster` (distinct hosts).
409fn synth_from_host(kind: EndpointKind, get_env: &dyn Fn(&str) -> Option<String>) -> Option<String> {
410    let host = nonempty(get_env("NEWT_DGX_HOST"))?;
411    let (port_var, default_port) = match kind {
412        EndpointKind::Ollama => ("NEWT_DGX_OLLAMA_PORT", "11434"),
413        EndpointKind::Vllm => ("NEWT_DGX_VLLM_PORT", "8000"),
414        EndpointKind::OllamaLb | EndpointKind::InCluster => return None,
415    };
416    let scheme = nonempty(get_env("NEWT_DGX_SCHEME")).unwrap_or_else(|| "http".to_string());
417    let port = nonempty(get_env(port_var)).unwrap_or_else(|| default_port.to_string());
418    Some(format!("{scheme}://{host}:{port}"))
419}
420
421// ---------------------------------------------------------------------------
422// Tests
423// ---------------------------------------------------------------------------
424
425#[cfg(test)]
426mod tests {
427    use super::*;
428
429    /// Build an injectable env lookup from a fixed set of pairs.
430    fn env_of<'a>(pairs: &'a [(&'a str, &'a str)]) -> impl Fn(&str) -> Option<String> + 'a {
431        move |key| {
432            pairs
433                .iter()
434                .find(|(k, _)| *k == key)
435                .map(|(_, v)| (*v).to_string())
436        }
437    }
438
439    // --- EndpointKind --------------------------------------------------
440
441    #[test]
442    fn endpoint_kind_str_roundtrip() {
443        for kind in EndpointKind::ALL {
444            let s = kind.as_str();
445            assert_eq!(s.parse::<EndpointKind>().unwrap(), kind);
446            assert_eq!(kind.to_string(), s);
447        }
448    }
449
450    #[test]
451    fn endpoint_kind_from_str_aliases() {
452        assert_eq!(
453            "LB".parse::<EndpointKind>().unwrap(),
454            EndpointKind::OllamaLb
455        );
456        assert_eq!(
457            "in-cluster".parse::<EndpointKind>().unwrap(),
458            EndpointKind::InCluster
459        );
460        assert_eq!(
461            "proxy".parse::<EndpointKind>().unwrap(),
462            EndpointKind::InCluster
463        );
464        assert_eq!(
465            "  Ollama ".parse::<EndpointKind>().unwrap(),
466            EndpointKind::Ollama
467        );
468    }
469
470    #[test]
471    fn endpoint_kind_from_str_rejects_unknown() {
472        let err = "bogus".parse::<EndpointKind>().unwrap_err();
473        assert!(
474            err.contains("bogus"),
475            "error should name the bad input: {err}"
476        );
477    }
478
479    #[test]
480    fn endpoint_kind_default_is_ollama() {
481        assert_eq!(EndpointKind::default(), EndpointKind::Ollama);
482    }
483
484    #[test]
485    fn endpoint_kind_openai_compat_only_vllm() {
486        assert!(EndpointKind::Vllm.is_openai_compatible());
487        assert!(!EndpointKind::Ollama.is_openai_compatible());
488        assert!(!EndpointKind::InCluster.is_openai_compatible());
489    }
490
491    #[test]
492    fn endpoint_kind_serde_snake_case() {
493        let json = serde_json::to_string(&EndpointKind::InCluster).unwrap();
494        assert_eq!(json, "\"in_cluster\"");
495        let back: EndpointKind = serde_json::from_str("\"ollama_lb\"").unwrap();
496        assert_eq!(back, EndpointKind::OllamaLb);
497    }
498
499    // --- node accessor -------------------------------------------------
500
501    #[test]
502    fn node_endpoint_accessor() {
503        let node = DgxNode {
504            name: "n".into(),
505            ollama: Some("o".into()),
506            vllm: Some("v".into()),
507            ..Default::default()
508        };
509        assert_eq!(node.endpoint(EndpointKind::Ollama), Some("o"));
510        assert_eq!(node.endpoint(EndpointKind::Vllm), Some("v"));
511        assert_eq!(node.endpoint(EndpointKind::OllamaLb), None);
512        assert_eq!(node.endpoint(EndpointKind::InCluster), None);
513    }
514
515    // --- active_node ---------------------------------------------------
516
517    #[test]
518    fn active_node_single_is_implicit() {
519        let cfg = DgxConfig {
520            nodes: vec![DgxNode {
521                name: "solo".into(),
522                ..Default::default()
523            }],
524            ..Default::default()
525        };
526        assert_eq!(cfg.active_node().unwrap().name, "solo");
527    }
528
529    #[test]
530    fn active_node_explicit_selection() {
531        let cfg = DgxConfig {
532            active_node: Some("b".into()),
533            nodes: vec![
534                DgxNode {
535                    name: "a".into(),
536                    ..Default::default()
537                },
538                DgxNode {
539                    name: "b".into(),
540                    ..Default::default()
541                },
542            ],
543            ..Default::default()
544        };
545        assert_eq!(cfg.active_node().unwrap().name, "b");
546    }
547
548    #[test]
549    fn active_node_unknown_name_errors() {
550        let cfg = DgxConfig {
551            active_node: Some("ghost".into()),
552            nodes: vec![DgxNode {
553                name: "real".into(),
554                ..Default::default()
555            }],
556            ..Default::default()
557        };
558        assert_eq!(
559            cfg.active_node().unwrap_err(),
560            DgxNotConfigured::NodeNotFound {
561                name: "ghost".into()
562            }
563        );
564    }
565
566    #[test]
567    fn active_node_zero_nodes_errors() {
568        let cfg = DgxConfig::default();
569        assert_eq!(cfg.active_node().unwrap_err(), DgxNotConfigured::NoNodes);
570    }
571
572    #[test]
573    fn active_node_many_without_active_errors() {
574        let cfg = DgxConfig {
575            nodes: vec![
576                DgxNode {
577                    name: "a".into(),
578                    ..Default::default()
579                },
580                DgxNode {
581                    name: "b".into(),
582                    ..Default::default()
583                },
584            ],
585            ..Default::default()
586        };
587        assert_eq!(
588            cfg.active_node().unwrap_err(),
589            DgxNotConfigured::NoActiveNode { count: 2 }
590        );
591    }
592
593    // --- endpoint resolution -------------------------------------------
594
595    #[test]
596    fn resolve_endpoint_uses_config_url() {
597        let cfg = DgxConfig::home_template();
598        assert_eq!(
599            cfg.resolve_endpoint_for_with(EndpointKind::Ollama, &env_of(&[]))
600                .unwrap(),
601            "https://REDACTED-HOST"
602        );
603        assert_eq!(
604            cfg.resolve_endpoint_for_with(EndpointKind::InCluster, &env_of(&[]))
605                .unwrap(),
606            "http://ollama-proxy.inference.svc.cluster.local:11434"
607        );
608    }
609
610    #[test]
611    fn resolve_active_endpoint_default_is_ollama() {
612        let cfg = DgxConfig::home_template();
613        assert_eq!(
614            cfg.resolve_endpoint_for_with(cfg.active_endpoint, &env_of(&[]))
615                .unwrap(),
616            "https://REDACTED-HOST"
617        );
618    }
619
620    #[test]
621    fn resolve_endpoint_env_url_wins() {
622        let cfg = DgxConfig::home_template();
623        let env = env_of(&[("NEWT_DGX_OLLAMA_URL", "http://override:1234")]);
624        assert_eq!(
625            cfg.resolve_endpoint_for_with(EndpointKind::Ollama, &env)
626                .unwrap(),
627            "http://override:1234"
628        );
629    }
630
631    #[test]
632    fn resolve_endpoint_empty_env_url_ignored() {
633        let cfg = DgxConfig::home_template();
634        let env = env_of(&[("NEWT_DGX_OLLAMA_URL", "   ")]);
635        assert_eq!(
636            cfg.resolve_endpoint_for_with(EndpointKind::Ollama, &env)
637                .unwrap(),
638            "https://REDACTED-HOST"
639        );
640    }
641
642    #[test]
643    fn resolve_endpoint_host_synthesis_ollama_default_port() {
644        let cfg = DgxConfig::default();
645        let env = env_of(&[("NEWT_DGX_HOST", "dgx.example")]);
646        assert_eq!(
647            cfg.resolve_endpoint_for_with(EndpointKind::Ollama, &env)
648                .unwrap(),
649            "http://dgx.example:11434"
650        );
651    }
652
653    #[test]
654    fn resolve_endpoint_host_synthesis_vllm_scheme_and_port() {
655        let cfg = DgxConfig::default();
656        let env = env_of(&[
657            ("NEWT_DGX_HOST", "dgx.example"),
658            ("NEWT_DGX_SCHEME", "https"),
659            ("NEWT_DGX_VLLM_PORT", "8443"),
660        ]);
661        assert_eq!(
662            cfg.resolve_endpoint_for_with(EndpointKind::Vllm, &env)
663                .unwrap(),
664            "https://dgx.example:8443"
665        );
666    }
667
668    #[test]
669    fn resolve_endpoint_host_synthesis_skips_lb_and_cluster() {
670        let cfg = DgxConfig::default();
671        let env = env_of(&[("NEWT_DGX_HOST", "dgx.example")]);
672        assert_eq!(
673            cfg.resolve_endpoint_for_with(EndpointKind::OllamaLb, &env)
674                .unwrap_err(),
675            DgxNotConfigured::NoNodes
676        );
677        assert_eq!(
678            cfg.resolve_endpoint_for_with(EndpointKind::InCluster, &env)
679                .unwrap_err(),
680            DgxNotConfigured::NoNodes
681        );
682    }
683
684    #[test]
685    fn resolve_endpoint_unset_on_node_errors() {
686        let cfg = DgxConfig {
687            nodes: vec![DgxNode {
688                name: "n".into(),
689                ollama: Some("o".into()),
690                ..Default::default()
691            }],
692            ..Default::default()
693        };
694        assert_eq!(
695            cfg.resolve_endpoint_for_with(EndpointKind::Vllm, &env_of(&[]))
696                .unwrap_err(),
697            DgxNotConfigured::EndpointUnset {
698                node: "n".into(),
699                kind: EndpointKind::Vllm
700            }
701        );
702    }
703
704    #[test]
705    fn resolve_endpoint_no_config_no_env_errors() {
706        let cfg = DgxConfig::default();
707        assert_eq!(
708            cfg.resolve_endpoint_for_with(EndpointKind::Ollama, &env_of(&[]))
709                .unwrap_err(),
710            DgxNotConfigured::NoNodes
711        );
712    }
713
714    // --- model / ssh ---------------------------------------------------
715
716    #[test]
717    fn resolve_active_model_config_then_env() {
718        let cfg = DgxConfig::home_template();
719        assert_eq!(
720            cfg.resolve_active_model_with(&env_of(&[])).unwrap(),
721            "qwen2.5-coder:32b"
722        );
723        let env = env_of(&[("NEWT_DGX_MODEL", "llama3.1:70b")]);
724        assert_eq!(cfg.resolve_active_model_with(&env).unwrap(), "llama3.1:70b");
725    }
726
727    #[test]
728    fn resolve_active_model_missing_errors() {
729        let cfg = DgxConfig::default();
730        assert_eq!(
731            cfg.resolve_active_model_with(&env_of(&[])).unwrap_err(),
732            DgxNotConfigured::NoActiveModel
733        );
734    }
735
736    #[test]
737    fn ssh_host_config_then_env() {
738        let cfg = DgxConfig::home_template();
739        assert_eq!(cfg.ssh_host_with(&env_of(&[])).unwrap(), "REDACTED-HOST");
740        let env = env_of(&[("NEWT_DGX_SSH_HOST", "other.host")]);
741        assert_eq!(cfg.ssh_host_with(&env).unwrap(), "other.host");
742    }
743
744    #[test]
745    fn ssh_host_missing_errors() {
746        let cfg = DgxConfig {
747            nodes: vec![DgxNode {
748                name: "n".into(),
749                ..Default::default()
750            }],
751            ..Default::default()
752        };
753        assert_eq!(
754            cfg.ssh_host_with(&env_of(&[])).unwrap_err(),
755            DgxNotConfigured::NoSshHost { node: "n".into() }
756        );
757    }
758
759    #[test]
760    fn ssh_user_precedence_chain() {
761        let cfg = DgxConfig {
762            nodes: vec![DgxNode {
763                name: "n".into(),
764                ssh_user: Some("nodeuser".into()),
765                ..Default::default()
766            }],
767            ..Default::default()
768        };
769        assert_eq!(
770            cfg.ssh_user_with(&env_of(&[("NEWT_DGX_SSH_USER", "envuser")])),
771            "envuser"
772        );
773        assert_eq!(cfg.ssh_user_with(&env_of(&[])), "nodeuser");
774
775        let bare = DgxConfig {
776            nodes: vec![DgxNode {
777                name: "n".into(),
778                ..Default::default()
779            }],
780            ..Default::default()
781        };
782        assert_eq!(bare.ssh_user_with(&env_of(&[("USER", "shell")])), "shell");
783        assert_eq!(bare.ssh_user_with(&env_of(&[])), "dgx");
784    }
785
786    // --- home template + round-trip ------------------------------------
787
788    #[test]
789    fn home_template_resolves() {
790        let cfg = DgxConfig::home_template();
791        assert_eq!(cfg.active_node().unwrap().name, "home");
792        let url = cfg
793            .resolve_endpoint_for_with(EndpointKind::Ollama, &env_of(&[]))
794            .unwrap();
795        assert!(
796            url.starts_with("https://"),
797            "direct DGX should be https: {url}"
798        );
799        assert!(
800            !url.contains(":11434"),
801            "Traefik endpoint should not carry :11434: {url}"
802        );
803        assert!(cfg.formation("coding").is_some());
804        assert!(cfg.formation("missing").is_none());
805    }
806
807    #[test]
808    fn dgx_config_toml_roundtrip() {
809        let cfg = DgxConfig::home_template();
810        let text = toml::to_string_pretty(&cfg).unwrap();
811        let back: DgxConfig = toml::from_str(&text).unwrap();
812        assert_eq!(cfg, back);
813    }
814
815    #[test]
816    fn real_env_wrappers_smoke() {
817        // Exercise the public, real-environment wrappers. home_template's
818        // values come from config, so these succeed regardless of the ambient
819        // environment (an env override would only change the value).
820        let cfg = DgxConfig::home_template();
821        assert!(cfg.resolve_endpoint().is_ok());
822        assert!(cfg.resolve_endpoint_for(EndpointKind::InCluster).is_ok());
823        assert!(cfg.resolve_active_model().is_ok());
824        assert!(cfg.ssh_host().is_ok());
825        assert!(!cfg.ssh_user().is_empty());
826    }
827}
newt_core/dgx.rs

newt_core/
dgx.rs
