falsegreen_agent/
server.rs

1//! Owned, loopback-only llama-server lifecycle management.
2
3use std::collections::VecDeque;
4use std::io::Read;
5use std::net::{Ipv4Addr, TcpListener};
6use std::path::{Path, PathBuf};
7use std::process::{Child, Command, ExitStatus, Stdio};
8use std::sync::{Arc, Mutex};
9use std::thread::{self, JoinHandle};
10use std::time::{Duration, Instant};
11
12use serde::{Deserialize, Serialize};
13use serde_json::Value;
14use thiserror::Error;
15
16use crate::hardware::RuntimeBackend;
17
18const START_ATTEMPTS: usize = 5;
19const DIAGNOSTIC_CAPACITY: usize = 128 * 1024;
20
21#[derive(Debug, Clone, PartialEq, Eq)]
22pub struct ServerLaunchConfig {
23    pub executable: PathBuf,
24    pub model: PathBuf,
25    pub model_identifier: String,
26    pub backend: RuntimeBackend,
27    pub context_tokens: u32,
28    pub logical_cpus: usize,
29    pub startup_timeout: Duration,
30    pub health_timeout: Duration,
31}
32
33impl ServerLaunchConfig {
34    pub fn validate(&self) -> Result<(), ServerError> {
35        if !self.executable.is_file() {
36            return Err(ServerError::InvalidConfig(format!(
37                "server executable is missing: {}",
38                self.executable.display()
39            )));
40        }
41        if !self.model.is_file() {
42            return Err(ServerError::InvalidConfig(format!(
43                "model artifact is missing: {}",
44                self.model.display()
45            )));
46        }
47        if self.model_identifier.is_empty() || self.model_identifier.chars().any(char::is_control) {
48            return Err(ServerError::InvalidConfig(
49                "model identifier must be non-empty and contain no control characters".to_owned(),
50            ));
51        }
52        if self.context_tokens == 0 {
53            return Err(ServerError::InvalidConfig(
54                "context size must be nonzero".to_owned(),
55            ));
56        }
57        if self.logical_cpus == 0 {
58            return Err(ServerError::InvalidConfig(
59                "logical CPU count must be nonzero".to_owned(),
60            ));
61        }
62        if self.startup_timeout.is_zero() || self.health_timeout.is_zero() {
63            return Err(ServerError::InvalidConfig(
64                "startup and health timeouts must be nonzero".to_owned(),
65            ));
66        }
67        Ok(())
68    }
69
70    #[must_use]
71    pub fn arguments(&self, port: u16) -> Vec<String> {
72        let threads = self.logical_cpus.clamp(1, 16).to_string();
73        let gpu_layers = if self.backend == RuntimeBackend::Cpu {
74            "0"
75        } else {
76            "999"
77        };
78        vec![
79            "--host".to_owned(),
80            Ipv4Addr::LOCALHOST.to_string(),
81            "--port".to_owned(),
82            port.to_string(),
83            "--model".to_owned(),
84            self.model.display().to_string(),
85            "--alias".to_owned(),
86            self.model_identifier.clone(),
87            "--ctx-size".to_owned(),
88            self.context_tokens.to_string(),
89            "--parallel".to_owned(),
90            "1".to_owned(),
91            "--threads".to_owned(),
92            threads.clone(),
93            "--threads-batch".to_owned(),
94            threads,
95            "--n-gpu-layers".to_owned(),
96            gpu_layers.to_owned(),
97            "--flash-attn".to_owned(),
98            "on".to_owned(),
99            "--jinja".to_owned(),
100            "--temp".to_owned(),
101            "0".to_owned(),
102            "--spec-type".to_owned(),
103            "none".to_owned(),
104            "--no-webui".to_owned(),
105        ]
106    }
107}
108
109#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
110pub struct ManagedServerEndpoint {
111    pub endpoint: String,
112    pub port: u16,
113    pub pid: u32,
114    pub reused: bool,
115}
116
117#[derive(Debug, Default)]
118pub struct ServerManager {
119    active: Option<ManagedServer>,
120}
121
122impl ServerManager {
123    #[must_use]
124    pub const fn new() -> Self {
125        Self { active: None }
126    }
127
128    /// Reuse is limited to a still-live process this manager created with the exact same config.
129    /// An unrelated listener is never adopted, even if it answers the health endpoint.
130    pub fn ensure_running(
131        &mut self,
132        config: ServerLaunchConfig,
133    ) -> Result<ManagedServerEndpoint, ServerError> {
134        config.validate()?;
135        let reusable = if let Some(active) = &mut self.active {
136            active.config == config
137                && active.is_live()?
138                && matches!(
139                    probe_health(&active.endpoint, config.health_timeout),
140                    Ok(true)
141                )
142                && matches!(
143                    probe_model_identity(
144                        &active.endpoint,
145                        config.health_timeout,
146                        &config.model_identifier
147                    ),
148                    Ok(true)
149                )
150        } else {
151            false
152        };
153        if reusable {
154            let active = self.active.as_ref().expect("reusable server is active");
155            return Ok(ManagedServerEndpoint {
156                endpoint: active.endpoint.clone(),
157                port: active.port,
158                pid: active.child.id(),
159                reused: true,
160            });
161        }
162        self.shutdown()?;
163
164        let mut last_error = None;
165        for _ in 0..START_ATTEMPTS {
166            let port = available_loopback_port()?;
167            match ManagedServer::start(config.clone(), port) {
168                Ok(server) => {
169                    let endpoint = ManagedServerEndpoint {
170                        endpoint: server.endpoint.clone(),
171                        port,
172                        pid: server.child.id(),
173                        reused: false,
174                    };
175                    self.active = Some(server);
176                    return Ok(endpoint);
177                }
178                Err(error @ ServerError::Exited { .. }) => last_error = Some(error),
179                Err(error) => return Err(error),
180            }
181        }
182        Err(last_error.unwrap_or_else(|| {
183            ServerError::Spawn("no loopback start attempt completed".to_owned())
184        }))
185    }
186
187    pub fn shutdown(&mut self) -> Result<(), ServerError> {
188        if let Some(server) = self.active.take() {
189            server.shutdown()?;
190        }
191        Ok(())
192    }
193
194    #[must_use]
195    pub fn diagnostics(&self) -> String {
196        self.active
197            .as_ref()
198            .map_or_else(String::new, |server| server.diagnostics.snapshot())
199    }
200}
201
202impl Drop for ServerManager {
203    fn drop(&mut self) {
204        let _ = self.shutdown();
205    }
206}
207
208#[derive(Debug, Error)]
209pub enum ServerError {
210    #[error("invalid server configuration: {0}")]
211    InvalidConfig(String),
212    #[error("could not reserve a private loopback port: {0}")]
213    Port(std::io::Error),
214    #[error("could not start managed llama-server: {0}")]
215    Spawn(String),
216    #[error(
217        "managed llama-server exited during startup with {status}; diagnostics:\n{diagnostics}"
218    )]
219    Exited {
220        status: ExitStatus,
221        diagnostics: String,
222    },
223    #[error(
224        "managed llama-server was not ready after {seconds} seconds; diagnostics:\n{diagnostics}"
225    )]
226    ReadinessTimeout { seconds: u64, diagnostics: String },
227    #[error("managed llama-server health response was invalid: {0}")]
228    InvalidHealth(String),
229    #[error("could not inspect or stop managed llama-server: {0}")]
230    Process(std::io::Error),
231}
232
233#[derive(Debug)]
234struct ManagedServer {
235    config: ServerLaunchConfig,
236    endpoint: String,
237    port: u16,
238    child: Child,
239    diagnostics: BoundedDiagnostics,
240    readers: Vec<JoinHandle<()>>,
241}
242
243impl ManagedServer {
244    fn start(config: ServerLaunchConfig, port: u16) -> Result<Self, ServerError> {
245        let working_directory = config.executable.parent().ok_or_else(|| {
246            ServerError::InvalidConfig("server executable has no parent directory".to_owned())
247        })?;
248        let mut command = Command::new(&config.executable);
249        command
250            .args(config.arguments(port))
251            .current_dir(working_directory)
252            .stdin(Stdio::null())
253            .stdout(Stdio::piped())
254            .stderr(Stdio::piped());
255        configure_library_path(&mut command, working_directory);
256        let mut child = spawn_with_short_retry(&mut command)?;
257        let diagnostics = BoundedDiagnostics::new(DIAGNOSTIC_CAPACITY);
258        let mut readers = Vec::new();
259        if let Some(stdout) = child.stdout.take() {
260            readers.push(diagnostics.capture("stdout", stdout));
261        }
262        if let Some(stderr) = child.stderr.take() {
263            readers.push(diagnostics.capture("stderr", stderr));
264        }
265        let endpoint = format!("http://{}:{port}", Ipv4Addr::LOCALHOST);
266        let started = Instant::now();
267        loop {
268            if let Some(status) = child.try_wait().map_err(ServerError::Process)? {
269                join_readers(&mut readers);
270                return Err(ServerError::Exited {
271                    status,
272                    diagnostics: diagnostics.snapshot(),
273                });
274            }
275            match probe_health(&endpoint, config.health_timeout) {
276                Ok(true)
277                    if probe_model_identity(
278                        &endpoint,
279                        config.health_timeout,
280                        &config.model_identifier,
281                    )
282                    .unwrap_or(false) =>
283                {
284                    thread::sleep(Duration::from_millis(50));
285                    if child.try_wait().map_err(ServerError::Process)?.is_none() {
286                        break;
287                    }
288                }
289                Ok(false) | Err(ServerError::Process(_)) => {}
290                Ok(true) => {}
291                Err(error) => {
292                    let _ = child.kill();
293                    let _ = child.wait();
294                    join_readers(&mut readers);
295                    return Err(error);
296                }
297            }
298            if started.elapsed() >= config.startup_timeout {
299                let _ = child.kill();
300                let _ = child.wait();
301                join_readers(&mut readers);
302                return Err(ServerError::ReadinessTimeout {
303                    seconds: config.startup_timeout.as_secs(),
304                    diagnostics: diagnostics.snapshot(),
305                });
306            }
307            thread::sleep(Duration::from_millis(100));
308        }
309        Ok(Self {
310            config,
311            endpoint,
312            port,
313            child,
314            diagnostics,
315            readers,
316        })
317    }
318
319    fn is_live(&mut self) -> Result<bool, ServerError> {
320        self.child
321            .try_wait()
322            .map(|status| status.is_none())
323            .map_err(ServerError::Process)
324    }
325
326    fn shutdown(mut self) -> Result<(), ServerError> {
327        let result = if self
328            .child
329            .try_wait()
330            .map_err(ServerError::Process)?
331            .is_none()
332        {
333            self.child
334                .kill()
335                .and_then(|()| self.child.wait().map(|_| ()))
336                .map_err(ServerError::Process)
337        } else {
338            Ok(())
339        };
340        join_readers(&mut self.readers);
341        result
342    }
343}
344
345impl Drop for ManagedServer {
346    fn drop(&mut self) {
347        if self.child.try_wait().ok().flatten().is_none() {
348            let _ = self.child.kill();
349            let _ = self.child.wait();
350        }
351        join_readers(&mut self.readers);
352    }
353}
354
355fn available_loopback_port() -> Result<u16, ServerError> {
356    let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).map_err(ServerError::Port)?;
357    listener
358        .local_addr()
359        .map(|address| address.port())
360        .map_err(ServerError::Port)
361}
362
363fn spawn_with_short_retry(command: &mut Command) -> Result<Child, ServerError> {
364    let mut last_error = None;
365    for _ in 0..3 {
366        match command.spawn() {
367            Ok(child) => return Ok(child),
368            Err(error) if executable_temporarily_busy(&error) => {
369                last_error = Some(error);
370                thread::sleep(Duration::from_millis(25));
371            }
372            Err(error) => return Err(ServerError::Spawn(error.to_string())),
373        }
374    }
375    Err(ServerError::Spawn(last_error.map_or_else(
376        || "spawn retry failed".to_owned(),
377        |error| error.to_string(),
378    )))
379}
380
381#[cfg(unix)]
382fn executable_temporarily_busy(error: &std::io::Error) -> bool {
383    error.raw_os_error() == Some(26)
384}
385
386#[cfg(not(unix))]
387fn executable_temporarily_busy(_error: &std::io::Error) -> bool {
388    false
389}
390
391fn probe_health(endpoint: &str, timeout: Duration) -> Result<bool, ServerError> {
392    let value = probe_json(endpoint, "/health", timeout)?;
393    Ok(value
394        .as_ref()
395        .and_then(|value| value.get("status"))
396        .and_then(Value::as_str)
397        == Some("ok"))
398}
399
400fn probe_model_identity(
401    endpoint: &str,
402    timeout: Duration,
403    expected_model: &str,
404) -> Result<bool, ServerError> {
405    let value = probe_json(endpoint, "/v1/models", timeout)?;
406    Ok(value
407        .as_ref()
408        .and_then(|value| value.get("data"))
409        .and_then(Value::as_array)
410        .is_some_and(|models| {
411            models
412                .iter()
413                .any(|model| model.get("id").and_then(Value::as_str) == Some(expected_model))
414        }))
415}
416
417fn probe_json(endpoint: &str, path: &str, timeout: Duration) -> Result<Option<Value>, ServerError> {
418    let agent = ureq::AgentBuilder::new().timeout(timeout).build();
419    let response = match agent.get(&format!("{endpoint}{path}")).call() {
420        Ok(response) => response,
421        Err(ureq::Error::Status(503, _)) | Err(ureq::Error::Status(425, _)) => return Ok(None),
422        Err(ureq::Error::Status(status, _)) => {
423            return Err(ServerError::InvalidHealth(format!("HTTP {status}")));
424        }
425        Err(ureq::Error::Transport(error)) => {
426            return Err(ServerError::Process(std::io::Error::other(
427                error.to_string(),
428            )));
429        }
430    };
431    let body = read_bounded(response.into_reader(), 64 * 1024)
432        .map_err(|error| ServerError::InvalidHealth(error.to_string()))?;
433    let value: Value = serde_json::from_str(&body)
434        .map_err(|error| ServerError::InvalidHealth(error.to_string()))?;
435    Ok(Some(value))
436}
437
438fn read_bounded(mut reader: impl Read, maximum: usize) -> std::io::Result<String> {
439    let limit = u64::try_from(maximum + 1).expect("health response bound fits u64");
440    let mut bytes = Vec::new();
441    reader.by_ref().take(limit).read_to_end(&mut bytes)?;
442    if bytes.len() > maximum {
443        return Err(std::io::Error::other("health response exceeded size limit"));
444    }
445    String::from_utf8(bytes)
446        .map_err(|error| std::io::Error::new(std::io::ErrorKind::InvalidData, error))
447}
448
449#[cfg(target_os = "linux")]
450fn configure_library_path(command: &mut Command, runtime_directory: &Path) {
451    let value = std::env::var_os("LD_LIBRARY_PATH").map_or_else(
452        || runtime_directory.as_os_str().to_owned(),
453        |existing| {
454            let mut paths = vec![runtime_directory.to_path_buf()];
455            paths.extend(std::env::split_paths(&existing));
456            std::env::join_paths(paths).unwrap_or_else(|_| runtime_directory.as_os_str().to_owned())
457        },
458    );
459    command.env("LD_LIBRARY_PATH", value);
460}
461
462#[cfg(target_os = "macos")]
463fn configure_library_path(command: &mut Command, runtime_directory: &Path) {
464    command.env("DYLD_LIBRARY_PATH", runtime_directory);
465}
466
467#[cfg(not(any(target_os = "linux", target_os = "macos")))]
468fn configure_library_path(_command: &mut Command, _runtime_directory: &Path) {}
469
470#[derive(Debug, Clone)]
471struct BoundedDiagnostics {
472    bytes: Arc<Mutex<VecDeque<u8>>>,
473    capacity: usize,
474}
475
476impl BoundedDiagnostics {
477    fn new(capacity: usize) -> Self {
478        Self {
479            bytes: Arc::new(Mutex::new(VecDeque::with_capacity(capacity))),
480            capacity,
481        }
482    }
483
484    fn capture(
485        &self,
486        source: &'static str,
487        mut reader: impl Read + Send + 'static,
488    ) -> JoinHandle<()> {
489        let diagnostics = self.clone();
490        thread::spawn(move || {
491            let mut buffer = [0_u8; 4096];
492            loop {
493                match reader.read(&mut buffer) {
494                    Ok(0) | Err(_) => break,
495                    Ok(count) => {
496                        diagnostics.push(format!("[{source}] ").as_bytes());
497                        diagnostics.push(&buffer[..count]);
498                    }
499                }
500            }
501        })
502    }
503
504    fn push(&self, bytes: &[u8]) {
505        let mut stored = self
506            .bytes
507            .lock()
508            .unwrap_or_else(|poisoned| poisoned.into_inner());
509        stored.extend(bytes);
510        while stored.len() > self.capacity {
511            stored.pop_front();
512        }
513    }
514
515    fn snapshot(&self) -> String {
516        let stored = self
517            .bytes
518            .lock()
519            .unwrap_or_else(|poisoned| poisoned.into_inner());
520        let bytes: Vec<u8> = stored.iter().copied().collect();
521        String::from_utf8_lossy(&bytes).into_owned()
522    }
523}
524
525fn join_readers(readers: &mut Vec<JoinHandle<()>>) {
526    for reader in readers.drain(..) {
527        let _ = reader.join();
528    }
529}
530
531#[cfg(test)]
532mod tests {
533    use std::io::{Read, Write};
534    use std::net::TcpListener;
535    use std::thread;
536    use std::time::Duration;
537
538    use super::{RuntimeBackend, ServerLaunchConfig, ServerManager, probe_health};
539
540    #[test]
541    fn launch_arguments_are_private_bounded_and_non_speculative() {
542        let config = ServerLaunchConfig {
543            executable: "/runtime/llama-server".into(),
544            model: "/models/model.gguf".into(),
545            model_identifier: "neohorse".to_owned(),
546            backend: RuntimeBackend::Rocm,
547            context_tokens: 8192,
548            logical_cpus: 64,
549            startup_timeout: Duration::from_secs(1),
550            health_timeout: Duration::from_secs(1),
551        };
552        let arguments = config.arguments(32123);
553        assert!(
554            arguments
555                .windows(2)
556                .any(|pair| pair == ["--host", "127.0.0.1"])
557        );
558        assert!(arguments.windows(2).any(|pair| pair == ["--port", "32123"]));
559        assert!(arguments.windows(2).any(|pair| pair == ["--threads", "16"]));
560        assert!(
561            arguments
562                .windows(2)
563                .any(|pair| pair == ["--spec-type", "none"])
564        );
565        assert!(arguments.contains(&"--no-webui".to_owned()));
566    }
567
568    #[test]
569    fn readiness_requires_the_expected_health_document() {
570        let listener = TcpListener::bind("127.0.0.1:0").expect("listener");
571        let address = listener.local_addr().expect("address");
572        thread::spawn(move || {
573            let (mut stream, _) = listener.accept().expect("accept");
574            let mut request = [0_u8; 4096];
575            let _ = stream.read(&mut request);
576            let body = r#"{"status":"ok"}"#;
577            write!(
578                stream,
579                "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
580                body.len()
581            )
582            .expect("response");
583        });
584        assert!(
585            probe_health(&format!("http://{address}"), Duration::from_secs(1)).expect("health")
586        );
587    }
588
589    #[test]
590    fn configuration_rejects_missing_files_before_spawn() {
591        let config = ServerLaunchConfig {
592            executable: "/missing/llama-server".into(),
593            model: "/missing/model.gguf".into(),
594            model_identifier: "neohorse".to_owned(),
595            backend: RuntimeBackend::Cpu,
596            context_tokens: 8192,
597            logical_cpus: 1,
598            startup_timeout: Duration::from_secs(1),
599            health_timeout: Duration::from_secs(1),
600        };
601        assert!(config.validate().is_err());
602    }
603
604    #[cfg(unix)]
605    #[test]
606    fn owns_reuses_and_shuts_down_only_its_child() {
607        use std::os::unix::fs::PermissionsExt;
608
609        let directory = tempfile::tempdir().expect("tempdir");
610        let executable = directory.path().join("fake-llama-server");
611        std::fs::write(
612            &executable,
613            r#"#!/usr/bin/env python3
614import json
615import sys
616from http.server import BaseHTTPRequestHandler, HTTPServer
617port = int(sys.argv[sys.argv.index("--port") + 1])
618class Handler(BaseHTTPRequestHandler):
619    def do_GET(self):
620        if self.path == "/v1/models":
621            value = {"data": [{"id": "fixture"}]}
622        else:
623            value = {"status": "ok"}
624        body = json.dumps(value).encode()
625        self.send_response(200)
626        self.send_header("Content-Length", str(len(body)))
627        self.end_headers()
628        self.wfile.write(body)
629    def log_message(self, *args):
630        pass
631HTTPServer(("127.0.0.1", port), Handler).serve_forever()
632"#,
633        )
634        .expect("script");
635        let mut permissions = std::fs::metadata(&executable)
636            .expect("metadata")
637            .permissions();
638        permissions.set_mode(0o700);
639        std::fs::set_permissions(&executable, permissions).expect("permissions");
640        let model = directory.path().join("model.gguf");
641        std::fs::write(&model, b"fixture").expect("model");
642        let config = ServerLaunchConfig {
643            executable,
644            model,
645            model_identifier: "fixture".to_owned(),
646            backend: RuntimeBackend::Cpu,
647            context_tokens: 128,
648            logical_cpus: 2,
649            startup_timeout: Duration::from_secs(3),
650            health_timeout: Duration::from_millis(250),
651        };
652        let mut manager = ServerManager::new();
653        let first = manager.ensure_running(config.clone()).expect("start");
654        assert!(!first.reused);
655        let second = manager.ensure_running(config).expect("reuse");
656        assert!(second.reused);
657        assert_eq!(first.pid, second.pid);
658        assert_eq!(first.endpoint, second.endpoint);
659        manager.shutdown().expect("shutdown");
660        assert!(manager.diagnostics().is_empty());
661    }
662}
falsegreen_agent/server.rs

falsegreen_agent/
server.rs
