Skip to main content

frink_server/
cli.rs

1//! `frink-server`'s argument surface: the llama.cpp-shaped command line
2//! both front ends parse, and the environment it lowers to.
3//!
4//! Split out of `lib.rs` under this repo's "a new file beats a new
5//! section" rule, with no behaviour change: the struct, the
6//! multi-character short-option rewrite (`-ngl`, `-np`, `-hf`), the
7//! device/GPU-layer value types and `apply_cli_overrides` all moved
8//! verbatim, along with the tests that cover them.
9//!
10//! One rule this module holds to, because the repo has already paid for
11//! breaking it: **a flag that is accepted must reach the thing it
12//! names.** `frink bench --n-gpu-layers 0` was documented as forcing
13//! CPU and did not, because the backend was decided before the flag was
14//! read. Every override here runs in `apply_cli_overrides`, *before*
15//! the runtime and the model loader read the environment.
16
17use std::fmt;
18use std::net::{IpAddr, Ipv4Addr, SocketAddr};
19use std::path::PathBuf;
20use std::str::FromStr;
21
22use clap::{Parser, ValueEnum};
23
24// `PartialEq` so frink-cli's serve tests can assert that both front
25// ends parse a command line into the SAME arguments, rather than
26// asserting field by field and missing whichever one is added next.
27#[derive(Parser, Debug, PartialEq)]
28// No `version` here on purpose. This struct is both `frink-server`'s
29// own argv and the body of frink-cli's `serve` subcommand, and clap
30// gives an embedded subcommand its own `--version` derived from the
31// variant name: `frink serve --version` printed `frink-serve 0.10.0`,
32// naming a binary nobody ships. The front end's own `--version` is the
33// truth, and both report the same workspace version anyway.
34#[command(
35    name = "frink-server",
36    about = "OpenAI-compatible Frink inference server"
37)]
38pub struct ServerArgs {
39    /// Model path (GGUF file or Kimi checkpoint directory).
40    #[arg(short = 'm', long = "model", value_name = "FILE")]
41    model: Option<String>,
42
43    /// Hugging Face repo to serve, `user/repo[:QUANT]`, llama.cpp's
44    /// `-hf`.
45    ///
46    /// Downloads into the frink cache on first use and reuses it
47    /// after, so `-hf TheBloke/Mixtral-8x7B-Instruct-v0.1-GGUF:Q4_K_M`
48    /// is the whole command. The tag after the colon is a QUANT LABEL,
49    /// not a git revision, and it matches without regard to case.
50    #[arg(
51        long = "hf-repo",
52        visible_alias = "hf",
53        value_name = "REPO[:QUANT]",
54        conflicts_with = "model"
55    )]
56    hf_repo: Option<String>,
57
58    /// Exact filename inside `--hf-repo`, llama.cpp's `-hff`.
59    ///
60    /// For a repo whose quant labels do not disambiguate, or a file
61    /// whose name carries no quant at all.
62    #[arg(long = "hf-file", value_name = "FILE", requires = "hf_repo")]
63    hf_file: Option<String>,
64
65    /// Context size, llama.cpp's `-c`. Sets `FRINK_CB_MAX_CONTEXT`.
66    ///
67    /// Unset means the ceiling is derived at load from the weights and
68    /// the per-token KV against the device budget, capped at the
69    /// model's trained context, which is usually what you want.
70    #[arg(short = 'c', long = "ctx-size", value_name = "N")]
71    ctx_size: Option<usize>,
72
73    /// Require `Authorization: Bearer <key>`, llama.cpp's `--api-key`.
74    /// Sets `FRINK_API_KEY`, which also gates `/admin`.
75    #[arg(long = "api-key", value_name = "KEY")]
76    api_key: Option<String>,
77
78    /// Read the API key from a file, llama.cpp's `--api-key-file`.
79    ///
80    /// Preferred over `--api-key` on a shared host: an argument is
81    /// visible in `ps` to every user on the machine.
82    #[arg(long = "api-key-file", value_name = "PATH", conflicts_with = "api_key")]
83    api_key_file: Option<std::path::PathBuf>,
84
85    /// Name this model answers to in `/v1/models` and in responses,
86    /// llama.cpp's `--alias`. Sets `FRINK_MODEL_NAME`.
87    #[arg(long = "alias", visible_alias = "model-alias", value_name = "NAME")]
88    alias: Option<String>,
89
90    /// KV cache dtype, llama.cpp's `--cache-type-k`. Metal only; the
91    /// CPU and CUDA KV cache is the host `Vec<f32>`.
92    ///
93    /// Validated against the same vocabulary `frink run` uses
94    /// (`frink_models::ctk`): the two flags set the same variable, and
95    /// an engine that refuses a value on one binary while silently
96    /// serving f16 on the other is two answers to one question.
97    #[arg(
98        long = "ctk",
99        visible_alias = "cache-type-k",
100        value_name = "TYPE",
101        value_parser = frink_models::ctk::parse_value
102    )]
103    ctk: Option<String>,
104
105    /// Accepted and already the default: frink always compiles and
106    /// evaluates the GGUF's own `tokenizer.chat_template`. llama.cpp
107    /// needs `--jinja` to do that, so a command copied from there
108    /// carries it, and dying on an unknown flag would be a worse answer
109    /// than saying "yes, always".
110    #[arg(long = "jinja", default_value_t = false)]
111    jinja: bool,
112
113    /// Refused rather than ignored: frink has no
114    /// template-free/sniffing mode to fall back to. See `--jinja`.
115    #[arg(long = "no-jinja", default_value_t = false)]
116    no_jinja: bool,
117
118    /// Accepted; frink does no warm-up pass, so there is none to skip.
119    #[arg(long = "no-warmup", default_value_t = false)]
120    no_warmup: bool,
121
122    /// Accepted. Fused attention is a backend decision here, not a
123    /// request-time one: it is on wherever the Metal kernels support
124    /// the shape (`FRINK_METAL_ATTN`).
125    #[arg(long = "flash-attn", visible_alias = "fa", value_name = "MODE", num_args = 0..=1, default_missing_value = "auto")]
126    flash_attn: Option<String>,
127
128    /// IP address to listen on.
129    #[arg(long, value_name = "HOST")]
130    host: Option<IpAddr>,
131
132    /// Port to listen on. `0` asks the kernel for a free one; the
133    /// actually-bound address is then announced on stdout (see
134    /// [`announce_ready`]), which is how a supervising process is meant
135    /// to learn it.
136    #[arg(long, value_name = "PORT")]
137    port: Option<u16>,
138
139    /// CPU threads (sets FRINK_CPU_THREADS and RAYON_NUM_THREADS).
140    #[arg(short = 't', long = "threads", value_name = "N")]
141    threads: Option<usize>,
142
143    /// Device used for offloading (`none` disables GPU use).
144    #[arg(
145        long = "device",
146        visible_alias = "dev",
147        value_name = "DEVICE",
148        ignore_case = true
149    )]
150    device: Option<OffloadDevice>,
151
152    /// Print available offload devices and exit.
153    #[arg(long = "list-devices", default_value_t = false)]
154    pub(crate) list_devices: bool,
155
156    /// GPU layers: `0`, a positive number, `auto`, or `all`.
157    ///
158    /// Partial placement is not implemented yet; any value above zero
159    /// currently enables all supported operations on the selected backend.
160    #[arg(
161        long = "n-gpu-layers",
162        visible_aliases = ["gpu-layers", "ngl"],
163        value_name = "N"
164    )]
165    n_gpu_layers: Option<GpuLayers>,
166
167    /// MCP tool-server config JSON (stub: listed in `/v1/models` metadata).
168    #[arg(long = "mcp-config", value_name = "PATH")]
169    pub(crate) mcp_config: Option<PathBuf>,
170
171    /// Exit when stdin reaches EOF (for a supervising parent process).
172    ///
173    /// Opt-in on purpose: a server started with stdin redirected from
174    /// `/dev/null` -- systemd, cron, `nohup` -- sees EOF immediately,
175    /// and making this the default would turn those into a server that
176    /// exits the moment it starts. A parent that *wants* the guarantee
177    /// (the desktop shell) passes the flag and keeps the pipe open.
178    #[arg(long = "exit-on-stdin-close", default_value_t = false)]
179    pub(crate) exit_on_stdin_close: bool,
180
181    /// Share one batched decode worker across concurrent requests
182    /// (llama.cpp `-cb`). Also sets `FRINK_CONTINUOUS_BATCHING=1`.
183    #[arg(
184        long = "cont-batching",
185        visible_aliases = ["continuous-batching", "cb"],
186        default_value_t = false
187    )]
188    cont_batching: bool,
189
190    /// Disable auto continuous batching on Metal
191    /// (`FRINK_CONTINUOUS_BATCHING=0`).
192    #[arg(
193        long = "no-cont-batching",
194        default_value_t = false,
195        conflicts_with = "cont_batching"
196    )]
197    no_cont_batching: bool,
198
199    /// Max concurrent sequences under continuous batching (llama.cpp
200    /// `-np`). Sets `FRINK_CB_MAX_SEQS`; implies `--cont-batching`
201    /// unless `--no-cont-batching` is set.
202    #[arg(long = "parallel", visible_alias = "np", value_name = "N")]
203    parallel: Option<usize>,
204
205    /// Logical maximum prompt tokens per forward pass, llama.cpp's
206    /// `-b`. See [`crate::prefill_batch`] for how it and `-ub` resolve
207    /// to the one number frink keeps.
208    #[arg(long = "batch-size", visible_alias = "b", value_name = "N")]
209    batch_size: Option<usize>,
210
211    /// Physical maximum prompt tokens per forward pass, llama.cpp's
212    /// `-ub`. Clamped to `--batch-size` when both are given.
213    #[arg(long = "ubatch-size", visible_alias = "ub", value_name = "N")]
214    ubatch_size: Option<usize>,
215
216    /// Directory slot files are saved into and restored from,
217    /// llama.cpp's `--slot-save-path`. Sets `FRINK_SLOT_SAVE_PATH`.
218    ///
219    /// `POST /slots/{id_slot}?action=save|restore` refuses with a 501
220    /// naming this flag while it is unset, exactly as llama.cpp does
221    /// (`tools/server/server-context.cpp:4538`): a server that would
222    /// write KV state to disk should have been told where.
223    #[arg(long = "slot-save-path", value_name = "DIR")]
224    slot_save_path: Option<PathBuf>,
225
226    /// LoRA adapter GGUF (llama.cpp's `--lora`), applied at scale 1.
227    /// Repeatable; comma-separated values are accepted as upstream
228    /// accepts them. Sets `FRINK_LORA` together with `--lora-scaled`,
229    /// and every load -- the first and each `/admin/models/load` --
230    /// attaches the same adapters or refuses the checkpoint by name.
231    #[arg(long = "lora", value_name = "FILE", action = clap::ArgAction::Append)]
232    lora: Vec<String>,
233
234    /// LoRA adapter with a scale, `FILE:SCALE` (llama.cpp's
235    /// `--lora-scaled`). Repeatable; adapters are numbered in the order
236    /// given, every `--lora` before every `--lora-scaled`, and that
237    /// number is the `id` `POST /lora-adapters` and a request's `lora`
238    /// field address.
239    #[arg(long = "lora-scaled", value_name = "FILE:SCALE", action = clap::ArgAction::Append)]
240    lora_scaled: Vec<String>,
241
242    /// Load the adapters but apply none of them until a
243    /// `POST /lora-adapters` sets their scales (llama.cpp's
244    /// `--lora-init-without-apply`). Sets `FRINK_LORA_INIT_WITHOUT_APPLY`.
245    #[arg(long = "lora-init-without-apply", default_value_t = false)]
246    lora_init_without_apply: bool,
247
248    /// Start even though another frink process is already holding a
249    /// model. Off by default: two models on one box do not share it,
250    /// they thrash it, and both serve slower than either would alone.
251    /// `FRINK_ALLOW_MULTIPLE_INSTANCES=1` does the same.
252    #[arg(long = "allow-multiple-instances", default_value_t = false)]
253    pub(crate) allow_multiple_instances: bool,
254
255    /// Token budget for thinking, llama.cpp's `--reasoning-budget`: -1
256    /// for unrestricted, 0 for immediate end, N>0 for a token budget
257    /// (default: -1). The server default a request's
258    /// `reasoning_budget_tokens` falls back to when it is absent or -1;
259    /// sets `FRINK_REASONING_BUDGET`. Enforced in the sampler: once N
260    /// tokens of thought have followed the opener, the closer is forced
261    /// so the answer still arrives.
262    #[arg(
263        long = "reasoning-budget",
264        value_name = "N",
265        allow_hyphen_values = true
266    )]
267    reasoning_budget: Option<i64>,
268
269    /// Continue a trailing assistant message instead of starting a new
270    /// turn, llama.cpp's `--prefill-assistant` (the default). A request
271    /// can still say `continue_final_message: false`.
272    #[arg(long = "prefill-assistant", default_value_t = false)]
273    prefill_assistant: bool,
274
275    /// Treat a trailing assistant message as a complete turn,
276    /// llama.cpp's `--no-prefill-assistant`. Sets
277    /// `FRINK_PREFILL_ASSISTANT=0`; a request's own
278    /// `continue_final_message` still wins.
279    #[arg(
280        long = "no-prefill-assistant",
281        default_value_t = false,
282        conflicts_with = "prefill_assistant"
283    )]
284    no_prefill_assistant: bool,
285}
286
287impl ServerArgs {
288    /// Parses `frink-server`'s own argv, including the llama.cpp-style
289    /// multi-character short options (`-ngl`, `-dev`) that clap cannot
290    /// express and which are rewritten to their long forms first.
291    ///
292    /// Public because frink-cli's `serve` subcommand hands the same
293    /// arguments to the same parser rather than reimplementing it.
294    pub fn parse_llama_style<I>(argv: I) -> Self
295    where
296        I: IntoIterator<Item = String>,
297    {
298        Self::parse_from(rewrite_llama_style_argv(argv.into_iter().collect()))
299    }
300}
301
302#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
303enum OffloadDevice {
304    Auto,
305    None,
306    Cpu,
307    Metal,
308    Cuda,
309}
310
311#[derive(Debug, Clone, Copy, PartialEq, Eq)]
312enum GpuLayers {
313    Auto,
314    All,
315    Count(u32),
316}
317
318impl GpuLayers {
319    fn offload_enabled(self) -> bool {
320        !matches!(self, Self::Count(0))
321    }
322}
323
324impl FromStr for GpuLayers {
325    type Err = String;
326
327    fn from_str(value: &str) -> Result<Self, Self::Err> {
328        match value {
329            "auto" => Ok(Self::Auto),
330            "all" => Ok(Self::All),
331            _ => value
332                .parse::<u32>()
333                .map(Self::Count)
334                .map_err(|_| "expected 0, a positive integer, 'auto', or 'all'".into()),
335        }
336    }
337}
338
339impl fmt::Display for GpuLayers {
340    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
341        match self {
342            Self::Auto => f.write_str("auto"),
343            Self::All => f.write_str("all"),
344            Self::Count(value) => value.fmt(f),
345        }
346    }
347}
348
349/// Whether this build of the server has the Metal kernels compiled in.
350///
351/// Exists for the front ends that link this library: frink-cli's
352/// `metal` feature has to forward into frink-server
353/// (`frink-server?/metal`) or `frink serve --device metal` refuses on
354/// a Metal host while `frink run` on the same binary uses it. That
355/// mismatch is one Cargo manifest edit away and compiles cleanly, so
356/// frink-cli asserts on this constant at compile time.
357pub const BUILT_WITH_METAL: bool = cfg!(feature = "metal");
358
359/// Whether this build of the server has the CUDA kernels compiled in.
360/// See [`BUILT_WITH_METAL`].
361pub const BUILT_WITH_CUDA: bool = cfg!(feature = "cuda");
362
363fn rewrite_llama_style_argv(args: Vec<String>) -> Vec<String> {
364    args.into_iter()
365        .map(|arg| match arg.as_str() {
366            "-ngl" => "--n-gpu-layers".into(),
367            "-dev" => "--device".into(),
368            "-cb" => "--cont-batching".into(),
369            "-np" => "--parallel".into(),
370            "-b" => "--batch-size".into(),
371            "-ub" => "--ubatch-size".into(),
372            // One token in llama.cpp's hand-written parser. clap sees
373            // `-h` followed by `f` and prints help, which is what
374            // `frink serve -hf repo:Q4_K_M` did: the flag looked
375            // absent rather than mis-spelled.
376            "-hf" => "--hf-repo".into(),
377            "-hff" => "--hf-file".into(),
378            _ => arg,
379        })
380        .collect()
381}
382
383fn cli_bind_addr(args: &ServerArgs, env_addr: Option<&str>) -> Option<String> {
384    if args.host.is_none() && args.port.is_none() {
385        return None;
386    }
387
388    let existing = env_addr.and_then(|value| value.parse::<SocketAddr>().ok());
389    let host = args
390        .host
391        .or_else(|| existing.map(|addr| addr.ip()))
392        .unwrap_or(IpAddr::V4(Ipv4Addr::LOCALHOST));
393    let port = args
394        .port
395        .or_else(|| existing.map(|addr| addr.port()))
396        .unwrap_or(8383);
397    Some(SocketAddr::new(host, port).to_string())
398}
399
400/// Resolves a `-hf` reference to a local path, downloading it once.
401///
402/// Progress goes to STDERR, not stdout: stdout carries the
403/// `frink.server.ready` line a supervising process parses, and a
404/// progress bar in the middle of it would break that contract.
405fn resolve_hf_repo(spec: &str, file: Option<&str>) -> anyhow::Result<String> {
406    let mut hf = frink_models::hub::HfRef::parse(spec);
407    if let Some(f) = file {
408        hf.file = Some(f.to_string());
409    }
410    eprintln!(
411        "frink: resolving {} on the Hub{}",
412        hf.repo,
413        hf.quant
414            .as_deref()
415            .map(|q| format!(" ({q})"))
416            .unwrap_or_default()
417    );
418
419    let mut last = std::time::Instant::now();
420    let mut draw = move |done: u64, total: Option<u64>| {
421        if last.elapsed() < std::time::Duration::from_millis(200) {
422            return;
423        }
424        last = std::time::Instant::now();
425        let mib = done as f64 / 1024.0 / 1024.0;
426        match total {
427            Some(t) if t > 0 => {
428                eprint!(
429                    "\r  {mib:>9.1} MiB  {:5.1}%",
430                    (done as f64 / t as f64) * 100.0
431                )
432            }
433            _ => eprint!("\r  {mib:>9.1} MiB"),
434        }
435    };
436
437    let (path, downloaded) = hf
438        .ensure_local(&mut draw)
439        .map_err(|e| anyhow::anyhow!("{e}"))?;
440    if downloaded {
441        eprintln!();
442        eprintln!("frink: downloaded {}", path.display());
443    } else {
444        eprintln!("frink: using cached {}", path.display());
445    }
446    Ok(path.to_string_lossy().into_owned())
447}
448
449pub(crate) fn apply_cli_overrides(args: &ServerArgs) -> anyhow::Result<()> {
450    if let Some(model) = &args.model {
451        // SAFETY: called before the runtime starts worker threads.
452        unsafe { std::env::set_var("FRINK_MODEL_PATH", model) };
453    }
454    if let Some(spec) = &args.hf_repo {
455        let path = resolve_hf_repo(spec, args.hf_file.as_deref())?;
456        // SAFETY: called before the runtime starts worker threads.
457        unsafe { std::env::set_var("FRINK_MODEL_PATH", &path) };
458    }
459    if let Some(n) = args.ctx_size {
460        if n == 0 {
461            anyhow::bail!("--ctx-size must be greater than zero");
462        }
463        // SAFETY: called before the runtime starts worker threads.
464        unsafe { std::env::set_var("FRINK_CB_MAX_CONTEXT", n.to_string()) };
465    }
466    if let Some(key) = &args.api_key {
467        // SAFETY: called before the runtime starts worker threads.
468        unsafe { std::env::set_var("FRINK_API_KEY", key) };
469    }
470    if let Some(path) = &args.api_key_file {
471        let key = std::fs::read_to_string(path)
472            .map_err(|e| anyhow::anyhow!("reading --api-key-file {}: {e}", path.display()))?;
473        let key = key.trim();
474        if key.is_empty() {
475            anyhow::bail!(
476                "--api-key-file {} is empty: an empty key would leave every route open, \
477                 which is the opposite of what passing the flag asked for",
478                path.display()
479            );
480        }
481        // SAFETY: called before the runtime starts worker threads.
482        unsafe { std::env::set_var("FRINK_API_KEY", key) };
483    }
484    if let Some(alias) = &args.alias {
485        // SAFETY: called before the runtime starts worker threads.
486        unsafe { std::env::set_var("FRINK_MODEL_NAME", alias) };
487    }
488    if let Some(ctk) = &args.ctk {
489        // SAFETY: called before the runtime starts worker threads.
490        unsafe { std::env::set_var("FRINK_CTK", ctk.trim()) };
491    }
492    // Refused by NAME rather than ignored. A prompt framed by a
493    // hand-written guess instead of the checkpoint's own template is
494    // the kind of wrong answer that reads as a model quality problem,
495    // so "frink cannot do that" is the honest reply.
496    if args.no_jinja {
497        anyhow::bail!(
498            "--no-jinja: frink has no template-free mode. It compiles and evaluates the GGUF's \
499             own tokenizer.chat_template, which is what llama.cpp's --jinja turns on, and there \
500             is no sniffing fallback to switch to. Use --no-cnv on `frink run` for a raw \
501             completion"
502        );
503    }
504    if let Some(mode) = &args.flash_attn {
505        let mode = mode.trim().to_ascii_lowercase();
506        if mode == "off" || mode == "disabled" || mode == "0" {
507            anyhow::bail!(
508                "--flash-attn off: fused attention is a backend property here, not a per-run \
509                 switch. Set FRINK_METAL_ATTN=0 to take the unfused Metal path, or --device cpu"
510            );
511        }
512    }
513
514    if let Some(addr) = cli_bind_addr(args, std::env::var("FRINK_ADDR").ok().as_deref()) {
515        // SAFETY: called before the runtime starts worker threads.
516        unsafe { std::env::set_var("FRINK_ADDR", addr) };
517    }
518
519    if let Some(threads) = args.threads {
520        if threads == 0 {
521            anyhow::bail!("--threads must be greater than zero");
522        }
523        // SAFETY: called before the runtime starts worker threads.
524        unsafe {
525            std::env::set_var("FRINK_CPU_THREADS", threads.to_string());
526            std::env::set_var("RAYON_NUM_THREADS", threads.to_string());
527        }
528    }
529
530    if args.device.is_none() && args.n_gpu_layers.is_none() {
531        // device overrides skipped
532    } else {
533        let layers = args.n_gpu_layers.unwrap_or(GpuLayers::Auto);
534        let device = if layers.offload_enabled() {
535            args.device.unwrap_or(OffloadDevice::Auto)
536        } else {
537            OffloadDevice::None
538        };
539
540        match device {
541            OffloadDevice::None | OffloadDevice::Cpu => unsafe {
542                std::env::set_var("FRINK_METAL", "0");
543                std::env::set_var("FRINK_METAL_ATTN", "0");
544                std::env::set_var("FRINK_CUDA", "0");
545            },
546            OffloadDevice::Auto => unsafe {
547                std::env::set_var("FRINK_METAL", "auto");
548                std::env::set_var("FRINK_CUDA", "auto");
549                if std::env::var_os("FRINK_METAL_ATTN").is_none() {
550                    std::env::set_var("FRINK_METAL_ATTN", "1");
551                }
552            },
553            OffloadDevice::Metal => {
554                #[cfg(not(feature = "metal"))]
555                {
556                    anyhow::bail!(
557                        "Metal requested but this binary was built without --features metal"
558                    );
559                }
560                #[cfg(feature = "metal")]
561                {
562                    if !frink_metal::MetalProfile::detect().available {
563                        anyhow::bail!("Metal requested but no Metal device is available");
564                    }
565                    unsafe {
566                        std::env::set_var("FRINK_METAL", "1");
567                        if std::env::var_os("FRINK_METAL_ATTN").is_none() {
568                            std::env::set_var("FRINK_METAL_ATTN", "1");
569                        }
570                        std::env::set_var("FRINK_CUDA", "0");
571                    }
572                }
573            }
574            OffloadDevice::Cuda => {
575                #[cfg(not(feature = "cuda"))]
576                {
577                    anyhow::bail!(
578                        "CUDA requested but this binary was built without --features cuda"
579                    );
580                }
581                #[cfg(feature = "cuda")]
582                {
583                    if !frink_cuda::HardwareProfile::detect().cuda_available {
584                        anyhow::bail!("CUDA requested but no CUDA device is available");
585                    }
586                    unsafe {
587                        std::env::set_var("FRINK_CUDA", "1");
588                        std::env::set_var("FRINK_METAL", "0");
589                        std::env::set_var("FRINK_METAL_ATTN", "0");
590                    }
591                }
592            }
593        }
594    }
595
596    if let Some(n) = args.parallel {
597        if n == 0 {
598            anyhow::bail!("--parallel must be greater than zero");
599        }
600        // SAFETY: called before the runtime starts worker threads.
601        unsafe { std::env::set_var("FRINK_CB_MAX_SEQS", n.to_string()) };
602    }
603
604    let lora_specs = frink_models::lora_attach::LoraSpec::from_flags(&args.lora, &args.lora_scaled)
605        .map_err(|e| anyhow::anyhow!("{e}"))?;
606    if !lora_specs.is_empty() {
607        for spec in &lora_specs {
608            if !spec.path.is_file() {
609                anyhow::bail!(
610                    "--lora {}: no such file. An adapter that cannot be opened would be \
611                     discovered at model load rather than at startup",
612                    spec.path.display()
613                );
614            }
615        }
616        let value = lora_specs
617            .iter()
618            .map(|s| format!("{}:{}", s.path.display(), s.scale))
619            .collect::<Vec<_>>()
620            .join(",");
621        // SAFETY: called before the runtime starts worker threads.
622        unsafe { std::env::set_var(crate::lora::ENV_SPECS, value) };
623    }
624    if args.lora_init_without_apply {
625        // SAFETY: called before the runtime starts worker threads.
626        unsafe { std::env::set_var(crate::lora::ENV_INIT_WITHOUT_APPLY, "1") };
627    }
628
629    if let Some(dir) = &args.slot_save_path {
630        if !dir.is_dir() {
631            anyhow::bail!(
632                "--slot-save-path {} is not a directory. Slots are written into it by name, so \
633                 a path that does not exist would be discovered on the first save rather than \
634                 at startup",
635                dir.display()
636            );
637        }
638        // SAFETY: called before the runtime starts worker threads.
639        unsafe { std::env::set_var("FRINK_SLOT_SAVE_PATH", dir) };
640    }
641
642    for (flag, value) in [
643        ("--batch-size", args.batch_size),
644        ("--ubatch-size", args.ubatch_size),
645    ] {
646        if value == Some(0) {
647            anyhow::bail!("{flag} must be greater than zero");
648        }
649    }
650    if let Some(chunk) = crate::prefill_batch::effective_chunk(args.batch_size, args.ubatch_size) {
651        // Both spellings, from one number and one array of names: the
652        // private decode loop and the batch scheduler each read their
653        // own variable, and an operator who names `-ub` must not have
654        // to know which path this server happens to be serving on.
655        for key in crate::prefill_batch::PREFILL_CHUNK_ENV_KEYS {
656            // SAFETY: called before the runtime starts worker threads.
657            unsafe { std::env::set_var(key, chunk.to_string()) };
658        }
659    }
660
661    if let Some(budget) = args.reasoning_budget {
662        // The same range check the request field applies, so the flag
663        // and the field cannot admit different values.
664        crate::reasoning_budget::BudgetTokens::parse(budget)
665            .map_err(|why| anyhow::anyhow!("--reasoning-budget: {why}"))?;
666        // SAFETY: called before the runtime starts worker threads.
667        unsafe {
668            std::env::set_var(
669                crate::reasoning_budget::SERVER_DEFAULT_ENV,
670                budget.to_string(),
671            )
672        };
673    }
674    if args.prefill_assistant {
675        // SAFETY: called before the runtime starts worker threads.
676        unsafe { std::env::set_var(crate::continuation::PREFILL_ASSISTANT_ENV, "1") };
677    } else if args.no_prefill_assistant {
678        // SAFETY: called before the runtime starts worker threads.
679        unsafe { std::env::set_var(crate::continuation::PREFILL_ASSISTANT_ENV, "0") };
680    }
681
682    if args.cont_batching {
683        // SAFETY: called before the runtime starts worker threads.
684        unsafe { std::env::set_var("FRINK_CONTINUOUS_BATCHING", "1") };
685    } else if args.no_cont_batching {
686        // SAFETY: called before the runtime starts worker threads.
687        unsafe { std::env::set_var("FRINK_CONTINUOUS_BATCHING", "0") };
688    } else if args.parallel.is_some() {
689        // llama.cpp `-np` is only meaningful with continuous batching.
690        // SAFETY: called before the runtime starts worker threads.
691        unsafe { std::env::set_var("FRINK_CONTINUOUS_BATCHING", "1") };
692    }
693
694    Ok(())
695}
696#[cfg(test)]
697mod tests {
698    use super::*;
699
700    /// The server and `frink run` set the same `FRINK_CTK` variable, so
701    /// they have to agree about which values exist. The server used to
702    /// accept anything and pass it through, which meant
703    /// `frink-server --ctk nonsense` served f16 without a word while
704    /// `frink run --ctk nonsense` refused.
705    #[test]
706    fn the_server_accepts_exactly_the_cache_types_the_cli_does() {
707        for good in ["f16", "q8_0", "q4_0", "fp8", "q5_1"] {
708            let a = ServerArgs::try_parse_from(["frink-server", "--ctk", good]);
709            assert!(a.is_ok(), "{good} was refused");
710        }
711        for bad in ["nonsense", "turbo4", "q3_k"] {
712            let e = ServerArgs::try_parse_from(["frink-server", "--ctk", bad])
713                .expect_err(&format!("{bad} was accepted"));
714            let msg = e.to_string();
715            assert!(msg.contains("unsupported cache type"), "{msg}");
716        }
717    }
718
719    #[test]
720    fn parses_llama_server_style_options() {
721        let argv = [
722            "frink-server",
723            "-m",
724            "model.gguf",
725            "--host",
726            "::1",
727            "--port",
728            "9000",
729            "-t",
730            "4",
731            "-dev",
732            "Metal",
733            "-ngl",
734            "all",
735        ]
736        .into_iter()
737        .map(String::from)
738        .collect();
739        let args = ServerArgs::try_parse_from(rewrite_llama_style_argv(argv)).unwrap();
740
741        assert_eq!(args.model.as_deref(), Some("model.gguf"));
742        assert_eq!(args.host, Some(IpAddr::V6(std::net::Ipv6Addr::LOCALHOST)));
743        assert_eq!(args.port, Some(9000));
744        assert_eq!(args.threads, Some(4));
745        assert_eq!(args.device, Some(OffloadDevice::Metal));
746        assert_eq!(args.n_gpu_layers, Some(GpuLayers::All));
747        assert_eq!(
748            cli_bind_addr(&args, Some("127.0.0.1:8383")).as_deref(),
749            Some("[::1]:9000")
750        );
751    }
752
753    #[test]
754    fn port_zero_survives_argument_parsing_as_a_real_request() {
755        // `--port 0` must reach the bind call intact: it is a request
756        // for a kernel-assigned port, not a missing value to default to
757        // 8383. The address it produces is deliberately provisional --
758        // the ready line reports what was actually bound.
759        let argv = ["frink-server", "--port", "0"]
760            .into_iter()
761            .map(String::from)
762            .collect();
763        let args = ServerArgs::try_parse_from(rewrite_llama_style_argv(argv)).unwrap();
764        assert_eq!(args.port, Some(0));
765        assert_eq!(
766            cli_bind_addr(&args, Some("127.0.0.1:8383")).as_deref(),
767            Some("127.0.0.1:0")
768        );
769    }
770
771    #[test]
772    fn parallel_flag_parses_and_rewrites_np() {
773        let argv = ["frink-server", "-np", "4"]
774            .into_iter()
775            .map(String::from)
776            .collect();
777        let args = ServerArgs::try_parse_from(rewrite_llama_style_argv(argv)).unwrap();
778        assert_eq!(args.parallel, Some(4));
779    }
780
781    /// `-b` and `-ub` are two tokens in llama.cpp's hand-written
782    /// parser and one token to clap, which sees `-b` as a short option
783    /// it has never heard of. The rewrite is what makes a copied
784    /// `llama-server ... -b 2048 -ub 512` command run here at all.
785    #[test]
786    fn batch_flags_parse_and_rewrite_their_llama_cpp_short_forms() {
787        let argv = ["frink-server", "-b", "2048", "-ub", "512"]
788            .into_iter()
789            .map(String::from)
790            .collect();
791        let args = ServerArgs::try_parse_from(rewrite_llama_style_argv(argv)).unwrap();
792        assert_eq!(args.batch_size, Some(2048));
793        assert_eq!(args.ubatch_size, Some(512));
794    }
795
796    /// Zero is not a batch size, and accepting it would make
797    /// `env_positive` panic the server later with a message naming an
798    /// environment variable the operator never set.
799    #[test]
800    fn a_zero_batch_size_is_refused_by_name_rather_than_lowered_to_the_environment() {
801        for flag in ["--batch-size", "--ubatch-size"] {
802            let args = ServerArgs::try_parse_from(
803                ["frink-server", flag, "0"].into_iter().map(String::from),
804            )
805            .unwrap();
806            let err = apply_cli_overrides(&args).unwrap_err().to_string();
807            assert!(err.contains(flag), "{flag}: {err}");
808        }
809    }
810
811    /// A path that is not a directory is refused at startup rather
812    /// than on the first save. The repo's rule about gates applies to
813    /// flags too: a `--slot-save-path` pointing at nothing looks
814    /// configured until somebody tries to use it, hours later.
815    #[test]
816    fn a_slot_save_path_that_is_not_a_directory_is_refused_at_startup() {
817        let args = ServerArgs::try_parse_from(
818            ["frink-server", "--slot-save-path", "/definitely/not/here"]
819                .into_iter()
820                .map(String::from),
821        )
822        .unwrap();
823        let err = apply_cli_overrides(&args).unwrap_err().to_string();
824        assert!(err.contains("--slot-save-path"), "{err}");
825        assert!(
826            std::env::var("FRINK_SLOT_SAVE_PATH").is_err(),
827            "a refused path must not have been lowered to the environment first"
828        );
829    }
830
831    /// llama.cpp's spelling and range (`common/arg.cpp:3608-3614`):
832    /// `-1`, `0` and `N` parse -- `-1` needs `allow_hyphen_values`, or
833    /// clap reads it as a flag -- and anything below `-1` is refused by
834    /// name before it is lowered to the environment.
835    #[test]
836    fn reasoning_budget_parses_llama_cpps_range_and_refuses_the_rest() {
837        for (value, expect) in [("-1", -1), ("0", 0), ("2000", 2000)] {
838            let args = ServerArgs::try_parse_from(
839                ["frink-server", "--reasoning-budget", value]
840                    .into_iter()
841                    .map(String::from),
842            )
843            .unwrap();
844            assert_eq!(args.reasoning_budget, Some(expect), "{value}");
845        }
846        let args = ServerArgs::try_parse_from(
847            ["frink-server", "--reasoning-budget", "-2"]
848                .into_iter()
849                .map(String::from),
850        )
851        .unwrap();
852        let err = apply_cli_overrides(&args).unwrap_err().to_string();
853        assert!(err.contains("--reasoning-budget"), "{err}");
854    }
855
856    /// Both spellings of llama.cpp's prefill switch parse, and they
857    /// conflict rather than letting the last one win silently.
858    #[test]
859    fn prefill_assistant_has_both_of_llama_cpps_spellings() {
860        let on = ServerArgs::try_parse_from(
861            ["frink-server", "--prefill-assistant"]
862                .into_iter()
863                .map(String::from),
864        )
865        .unwrap();
866        assert!(on.prefill_assistant && !on.no_prefill_assistant);
867        let off = ServerArgs::try_parse_from(
868            ["frink-server", "--no-prefill-assistant"]
869                .into_iter()
870                .map(String::from),
871        )
872        .unwrap();
873        assert!(off.no_prefill_assistant && !off.prefill_assistant);
874        assert!(ServerArgs::try_parse_from(
875            [
876                "frink-server",
877                "--prefill-assistant",
878                "--no-prefill-assistant"
879            ]
880            .into_iter()
881            .map(String::from),
882        )
883        .is_err());
884    }
885
886    #[test]
887    fn stdin_close_exit_is_opt_in() {
888        // Default off: a server whose stdin is /dev/null (systemd, cron,
889        // nohup) would otherwise exit the instant it started.
890        let args =
891            ServerArgs::try_parse_from(["frink-server"].into_iter().map(String::from)).unwrap();
892        assert!(!args.exit_on_stdin_close);
893        let args = ServerArgs::try_parse_from(
894            ["frink-server", "--exit-on-stdin-close"]
895                .into_iter()
896                .map(String::from),
897        )
898        .unwrap();
899        assert!(args.exit_on_stdin_close);
900    }
901}