Skip to main content

tsift_local_model/
lib.rs

1use anyhow::{Context, Result, bail};
2use serde::{Deserialize, Serialize};
3use std::collections::BTreeMap;
4use std::fs;
5use std::io::Write;
6use std::path::{Path, PathBuf};
7use std::process::Command;
8use std::time::{SystemTime, UNIX_EPOCH};
9
10pub const DEFAULT_DESKTOP_RUNTIME_MARGIN_MIB: u64 = 4096;
11pub const DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB: u64 = 768;
12pub const DEFAULT_IDLE_TTL_SECONDS: u64 = 0;
13
14/// Default llama.cpp router unload endpoint. Also the value llama-server
15/// listens on by default. Override via `--provider-endpoint` or the
16/// `TSIFT_LLAMA_CPP_ENDPOINT` env var when this port is taken (e.g. by a
17/// local WordPress instance at 8080).
18pub const DEFAULT_LLAMA_CPP_ENDPOINT: &str = "http://127.0.0.1:8080/models/unload";
19/// Default Ollama generate endpoint. Override via `--provider-endpoint` or
20/// the `TSIFT_OLLAMA_ENDPOINT` env var.
21pub const DEFAULT_OLLAMA_ENDPOINT: &str = "http://127.0.0.1:11434/api/generate";
22/// Default vLLM sleep endpoint. Override via `--provider-endpoint` or the
23/// `TSIFT_VLLM_ENDPOINT` env var.
24pub const DEFAULT_VLLM_ENDPOINT: &str = "http://127.0.0.1:8000/sleep";
25
26/// Env var override for the llama.cpp router unload endpoint.
27pub const LLAMA_CPP_ENDPOINT_ENV_VAR: &str = "TSIFT_LLAMA_CPP_ENDPOINT";
28/// Env var override for the Ollama generate endpoint.
29pub const OLLAMA_ENDPOINT_ENV_VAR: &str = "TSIFT_OLLAMA_ENDPOINT";
30/// Env var override for the vLLM sleep endpoint.
31pub const VLLM_ENDPOINT_ENV_VAR: &str = "TSIFT_VLLM_ENDPOINT";
32
33#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
34pub enum ModelRole {
35    Extract,
36    Embed,
37    Rerank,
38}
39
40#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
41pub enum ProviderKind {
42    LlamaCpp,
43    Ollama,
44    Vllm,
45    HashFallback,
46}
47
48#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
49pub enum UnloadStrategy {
50    ProcessExit,
51    OllamaKeepAliveZero,
52    LlamaCppRouterUnload,
53    VllmSleep,
54    None,
55}
56
57#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
58pub enum ConcurrencyClass {
59    ExclusiveLargeGpu,
60    SharedSmallGpu,
61    CpuOrHash,
62}
63
64#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
65pub enum LeaseMode {
66    Exclusive,
67    Shared,
68    CpuOrHash,
69}
70
71#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
72pub enum UnloadActionKind {
73    ProviderApi,
74    ProcessExit,
75    Sleep,
76    Noop,
77}
78
79#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
80pub struct ModelProfile {
81    pub id: &'static str,
82    pub label: &'static str,
83    pub provider: ProviderKind,
84    pub model_ref: &'static str,
85    pub quantization: &'static str,
86    pub roles: Vec<ModelRole>,
87    pub context_tokens: u32,
88    pub estimated_weights_mib: u64,
89    pub estimated_kv_mib: u64,
90    pub runtime_margin_mib: u64,
91    pub concurrency: ConcurrencyClass,
92    pub unload_strategy: UnloadStrategy,
93    pub notes: &'static str,
94}
95
96impl ModelProfile {
97    pub fn estimated_total_mib(&self) -> u64 {
98        self.estimated_weights_mib + self.estimated_kv_mib + self.runtime_margin_mib
99    }
100
101    pub fn supports_role(&self, role: &ModelRole) -> bool {
102        self.roles.iter().any(|candidate| candidate == role)
103    }
104}
105
106#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
107pub struct GpuProcess {
108    pub pid: Option<u32>,
109    pub process_name: String,
110    pub used_memory_mib: Option<u64>,
111}
112
113#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
114pub struct GpuProbe {
115    pub timestamp_unix_seconds: Option<u64>,
116    pub available: bool,
117    pub gpu_name: Option<String>,
118    pub driver_version: Option<String>,
119    pub total_vram_mib: Option<u64>,
120    pub used_vram_mib: Option<u64>,
121    pub free_vram_mib: Option<u64>,
122    pub processes: Vec<GpuProcess>,
123    pub error: Option<String>,
124}
125
126impl GpuProbe {
127    pub fn unavailable(error: impl Into<String>) -> Self {
128        Self {
129            timestamp_unix_seconds: Some(current_unix_seconds()),
130            available: false,
131            gpu_name: None,
132            driver_version: None,
133            total_vram_mib: None,
134            used_vram_mib: None,
135            free_vram_mib: None,
136            processes: Vec::new(),
137            error: Some(error.into()),
138        }
139    }
140
141    pub fn synthetic_vram(used_vram_mib: u64) -> Self {
142        Self {
143            timestamp_unix_seconds: Some(current_unix_seconds()),
144            available: true,
145            gpu_name: Some("synthetic GPU".to_string()),
146            driver_version: None,
147            total_vram_mib: None,
148            used_vram_mib: Some(used_vram_mib),
149            free_vram_mib: None,
150            processes: Vec::new(),
151            error: None,
152        }
153    }
154}
155
156#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
157pub struct ProfileSelection {
158    pub profile: ModelProfile,
159    pub selectable: bool,
160    pub reason: String,
161}
162
163#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
164pub struct LocalModelStatusReport {
165    pub gpu_probe: GpuProbe,
166    pub extractor_profiles: Vec<ProfileSelection>,
167    pub embedding_profiles: Vec<ProfileSelection>,
168    pub recommended_extractor: Option<String>,
169    pub recommended_embedding: Option<String>,
170    pub notes: Vec<String>,
171}
172
173#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
174pub struct ProviderUnloadAction {
175    pub kind: UnloadActionKind,
176    pub label: String,
177    pub command: Option<Vec<String>>,
178    pub http_method: Option<String>,
179    pub endpoint: Option<String>,
180    pub body_json: Option<String>,
181    pub required: bool,
182}
183
184#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
185pub struct LocalModelLease {
186    pub lease_id: String,
187    pub mode: LeaseMode,
188    pub profile: ModelProfile,
189    pub pre_load_gpu_probe: GpuProbe,
190    pub provider_endpoint: Option<String>,
191    pub provider_pid: Option<u32>,
192    pub idle_ttl_seconds: u64,
193    pub unload_strategy: UnloadStrategy,
194    pub unload_actions: Vec<ProviderUnloadAction>,
195}
196
197#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
198pub enum VramCleanupStatus {
199    Proven,
200    ProvenByExternalAccounting,
201    NotProven,
202    ProbeUnavailable,
203}
204
205#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
206pub struct VramCleanupEvaluation {
207    pub status: VramCleanupStatus,
208    pub cleanup_proven: bool,
209    pub pre_used_mib: Option<u64>,
210    pub post_used_mib: Option<u64>,
211    pub allowed_post_used_mib: Option<u64>,
212    pub used_delta_mib: Option<i64>,
213    pub external_process_delta_mib: u64,
214    pub blocking_processes: Vec<GpuProcess>,
215    pub reason: String,
216}
217
218#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
219pub struct LocalModelLifecycleReport {
220    pub lease: LocalModelLease,
221    pub post_unload_gpu_probe: GpuProbe,
222    pub cleanup: VramCleanupEvaluation,
223    pub notes: Vec<String>,
224}
225
226pub fn default_model_profiles() -> Vec<ModelProfile> {
227    vec![
228        ModelProfile {
229            id: "qwen3-32b-q4",
230            label: "Qwen3-32B 4-bit",
231            provider: ProviderKind::LlamaCpp,
232            model_ref: "Qwen/Qwen3-32B-GGUF",
233            quantization: "q4",
234            roles: vec![ModelRole::Extract],
235            context_tokens: 32_768,
236            estimated_weights_mib: 20_500,
237            estimated_kv_mib: 4_096,
238            runtime_margin_mib: DEFAULT_DESKTOP_RUNTIME_MARGIN_MIB,
239            concurrency: ConcurrencyClass::ExclusiveLargeGpu,
240            unload_strategy: UnloadStrategy::LlamaCppRouterUnload,
241            notes: "default quality extractor/reasoner for a clear RTX 5090",
242        },
243        ModelProfile {
244            id: "qwen3-30b-a3b-instruct-2507-q4",
245            label: "Qwen3-30B-A3B-Instruct-2507 4-bit",
246            provider: ProviderKind::LlamaCpp,
247            model_ref: "Qwen/Qwen3-30B-A3B-Instruct-2507",
248            quantization: "q4",
249            roles: vec![ModelRole::Extract],
250            context_tokens: 262_144,
251            estimated_weights_mib: 19_000,
252            estimated_kv_mib: 4_096,
253            runtime_margin_mib: DEFAULT_DESKTOP_RUNTIME_MARGIN_MIB,
254            concurrency: ConcurrencyClass::ExclusiveLargeGpu,
255            unload_strategy: UnloadStrategy::LlamaCppRouterUnload,
256            notes: "throughput and long-context extractor fallback",
257        },
258        ModelProfile {
259            id: "qwen3-embedding-0.6b",
260            label: "Qwen3-Embedding-0.6B",
261            provider: ProviderKind::LlamaCpp,
262            model_ref: "Qwen/Qwen3-Embedding-0.6B-GGUF",
263            quantization: "q8_or_f16",
264            roles: vec![ModelRole::Embed, ModelRole::Rerank],
265            context_tokens: 32_768,
266            estimated_weights_mib: 1_200,
267            estimated_kv_mib: 512,
268            runtime_margin_mib: 1_024,
269            concurrency: ConcurrencyClass::SharedSmallGpu,
270            unload_strategy: UnloadStrategy::LlamaCppRouterUnload,
271            notes: "default low-pressure embedding companion",
272        },
273        ModelProfile {
274            id: "qwen3-embedding-4b",
275            label: "Qwen3-Embedding-4B",
276            provider: ProviderKind::LlamaCpp,
277            model_ref: "Qwen/Qwen3-Embedding-4B",
278            quantization: "q4_or_q8",
279            roles: vec![ModelRole::Embed, ModelRole::Rerank],
280            context_tokens: 32_768,
281            estimated_weights_mib: 4_200,
282            estimated_kv_mib: 1_024,
283            runtime_margin_mib: 1_024,
284            concurrency: ConcurrencyClass::SharedSmallGpu,
285            unload_strategy: UnloadStrategy::LlamaCppRouterUnload,
286            notes: "higher-quality embedding candidate",
287        },
288        ModelProfile {
289            id: "qwen3-embedding-8b",
290            label: "Qwen3-Embedding-8B",
291            provider: ProviderKind::LlamaCpp,
292            model_ref: "Qwen/Qwen3-Embedding-8B",
293            quantization: "q4_or_q8",
294            roles: vec![ModelRole::Embed, ModelRole::Rerank],
295            context_tokens: 32_768,
296            estimated_weights_mib: 8_200,
297            estimated_kv_mib: 2_048,
298            runtime_margin_mib: 2_048,
299            concurrency: ConcurrencyClass::SharedSmallGpu,
300            unload_strategy: UnloadStrategy::LlamaCppRouterUnload,
301            notes: "benchmark when vector quality matters",
302        },
303        ModelProfile {
304            id: "qwen3.5-35b-a3b-q4",
305            label: "Qwen3.5-35B-A3B 4-bit",
306            provider: ProviderKind::LlamaCpp,
307            model_ref: "Qwen/Qwen3.5-35B-A3B",
308            quantization: "q4",
309            roles: vec![ModelRole::Extract],
310            context_tokens: 128_000,
311            estimated_weights_mib: 24_000,
312            estimated_kv_mib: 8_192,
313            runtime_margin_mib: DEFAULT_DESKTOP_RUNTIME_MARGIN_MIB,
314            concurrency: ConcurrencyClass::ExclusiveLargeGpu,
315            unload_strategy: UnloadStrategy::LlamaCppRouterUnload,
316            notes: "benchmark-only until a reduced-context single-5090 profile is proven",
317        },
318        ModelProfile {
319            id: "tsift-local-hash-v1",
320            label: "tsift local hash fallback",
321            provider: ProviderKind::HashFallback,
322            model_ref: "builtin",
323            quantization: "none",
324            roles: vec![ModelRole::Embed],
325            context_tokens: 0,
326            estimated_weights_mib: 0,
327            estimated_kv_mib: 0,
328            runtime_margin_mib: 0,
329            concurrency: ConcurrencyClass::CpuOrHash,
330            unload_strategy: UnloadStrategy::None,
331            notes: "deterministic fallback for tests and offline runs",
332        },
333        // ---- Ollama-native profiles (#lmlazy) ----
334        // model_ref matches an `ollama pull` tag so the unload lifecycle targets
335        // the live Ollama instance (TSIFT_OLLAMA_ENDPOINT / OLLAMA_HOST) instead
336        // of a llama.cpp router endpoint. Estimated VRAM mirrors the llama.cpp
337        // equivalents; Ollama loads the same GGUF weights onto the GPU.
338        ModelProfile {
339            id: "qwen3-32b-q4-ollama",
340            label: "Qwen3-32B 4-bit (Ollama)",
341            provider: ProviderKind::Ollama,
342            model_ref: "hf.co/Qwen/Qwen3-32B-GGUF:Q4_K_M",
343            quantization: "q4",
344            roles: vec![ModelRole::Extract],
345            context_tokens: 32_768,
346            estimated_weights_mib: 20_500,
347            estimated_kv_mib: 4_096,
348            runtime_margin_mib: DEFAULT_DESKTOP_RUNTIME_MARGIN_MIB,
349            concurrency: ConcurrencyClass::ExclusiveLargeGpu,
350            unload_strategy: UnloadStrategy::OllamaKeepAliveZero,
351            notes: "default Ollama-served quality extractor (lazy: keep_alive:0 unloads VRAM)",
352        },
353        ModelProfile {
354            id: "qwen3-embedding-0.6b-ollama",
355            label: "Qwen3-Embedding-0.6B (Ollama)",
356            provider: ProviderKind::Ollama,
357            model_ref: "hf.co/Qwen/Qwen3-Embedding-0.6B-GGUF",
358            quantization: "q8_or_f16",
359            roles: vec![ModelRole::Embed, ModelRole::Rerank],
360            context_tokens: 32_768,
361            estimated_weights_mib: 1_200,
362            estimated_kv_mib: 512,
363            runtime_margin_mib: 1_024,
364            concurrency: ConcurrencyClass::SharedSmallGpu,
365            unload_strategy: UnloadStrategy::OllamaKeepAliveZero,
366            notes: "default low-pressure Ollama-served embedding companion",
367        },
368    ]
369}
370
371pub fn probe_nvidia_smi() -> GpuProbe {
372    let output = Command::new("nvidia-smi")
373        .args([
374            "--query-gpu=name,driver_version,memory.total,memory.used,memory.free",
375            "--format=csv,noheader,nounits",
376        ])
377        .output();
378
379    let output = match output {
380        Ok(output) => output,
381        Err(error) => return GpuProbe::unavailable(format!("nvidia-smi unavailable: {error}")),
382    };
383
384    if !output.status.success() {
385        return GpuProbe::unavailable(format!(
386            "nvidia-smi failed: {}",
387            String::from_utf8_lossy(&output.stderr).trim()
388        ));
389    }
390
391    let stdout = String::from_utf8_lossy(&output.stdout);
392    match parse_gpu_query(stdout.lines().next().unwrap_or_default()) {
393        Ok(mut probe) => {
394            probe.processes = query_nvidia_compute_processes();
395            probe
396        }
397        Err(error) => GpuProbe::unavailable(error.to_string()),
398    }
399}
400
401pub fn build_status_report(probe_gpu: bool) -> LocalModelStatusReport {
402    let gpu_probe = if probe_gpu {
403        probe_nvidia_smi()
404    } else {
405        GpuProbe::unavailable("gpu probe skipped")
406    };
407    build_status_report_with_probe(gpu_probe)
408}
409
410pub fn build_status_report_with_probe(gpu_probe: GpuProbe) -> LocalModelStatusReport {
411    let profiles = default_model_profiles();
412    let extractor_profiles = rank_profiles_for_role(&profiles, &gpu_probe, ModelRole::Extract);
413    let embedding_profiles = rank_profiles_for_role(&profiles, &gpu_probe, ModelRole::Embed);
414    let recommended_extractor = extractor_profiles
415        .iter()
416        .find(|selection| selection.selectable)
417        .map(|selection| selection.profile.id.to_string());
418    let recommended_embedding = embedding_profiles
419        .iter()
420        .find(|selection| selection.selectable)
421        .map(|selection| selection.profile.id.to_string());
422
423    let mut notes = vec![
424        "large 30B/32B extractor profiles are single-lease on one RTX 5090".to_string(),
425        "use provider unload hooks or process exit after each batch to clear VRAM".to_string(),
426    ];
427    if !gpu_probe.available {
428        notes.push("GPU probe unavailable; profile fit is conservative".to_string());
429    }
430
431    LocalModelStatusReport {
432        gpu_probe,
433        extractor_profiles,
434        embedding_profiles,
435        recommended_extractor,
436        recommended_embedding,
437        notes,
438    }
439}
440
441pub fn profile_by_id(profile_id: &str) -> Option<ModelProfile> {
442    default_model_profiles()
443        .into_iter()
444        .find(|profile| profile.id == profile_id)
445}
446
447pub fn lease_mode_for_profile(profile: &ModelProfile) -> LeaseMode {
448    match profile.concurrency {
449        ConcurrencyClass::ExclusiveLargeGpu => LeaseMode::Exclusive,
450        ConcurrencyClass::SharedSmallGpu => LeaseMode::Shared,
451        ConcurrencyClass::CpuOrHash => LeaseMode::CpuOrHash,
452    }
453}
454
455pub fn build_local_model_lease(
456    profile: ModelProfile,
457    pre_load_gpu_probe: GpuProbe,
458    provider_endpoint: Option<String>,
459    provider_pid: Option<u32>,
460    idle_ttl_seconds: u64,
461) -> LocalModelLease {
462    let timestamp = pre_load_gpu_probe
463        .timestamp_unix_seconds
464        .unwrap_or_else(current_unix_seconds);
465    let lease_id = format!("{}-{timestamp}", profile.id);
466    let unload_actions = build_unload_actions(&profile, provider_endpoint.as_deref(), provider_pid);
467
468    LocalModelLease {
469        lease_id,
470        mode: lease_mode_for_profile(&profile),
471        unload_strategy: profile.unload_strategy.clone(),
472        profile,
473        pre_load_gpu_probe,
474        provider_endpoint,
475        provider_pid,
476        idle_ttl_seconds,
477        unload_actions,
478    }
479}
480
481pub fn build_unload_actions(
482    profile: &ModelProfile,
483    provider_endpoint: Option<&str>,
484    provider_pid: Option<u32>,
485) -> Vec<ProviderUnloadAction> {
486    match profile.unload_strategy {
487        UnloadStrategy::LlamaCppRouterUnload => {
488            let endpoint = resolve_provider_endpoint(&profile.unload_strategy, provider_endpoint);
489            let mut actions = vec![ProviderUnloadAction {
490                kind: UnloadActionKind::ProviderApi,
491                label: "llama.cpp router unload".to_string(),
492                command: None,
493                http_method: Some("POST".to_string()),
494                endpoint: Some(endpoint),
495                body_json: Some(format!(r#"{{"model":"{}"}}"#, profile.model_ref)),
496                required: true,
497            }];
498            if let Some(pid) = provider_pid {
499                actions.push(process_exit_action(
500                    pid,
501                    "terminate llama.cpp worker if unload is not proven",
502                ));
503            }
504            actions
505        }
506        UnloadStrategy::OllamaKeepAliveZero => vec![
507            ProviderUnloadAction {
508                kind: UnloadActionKind::ProviderApi,
509                label: "ollama keep_alive zero".to_string(),
510                command: None,
511                http_method: Some("POST".to_string()),
512                endpoint: Some(resolve_provider_endpoint(
513                    &profile.unload_strategy,
514                    provider_endpoint,
515                )),
516                body_json: Some(format!(
517                    r#"{{"model":"{}","prompt":"","keep_alive":0}}"#,
518                    profile.model_ref
519                )),
520                required: true,
521            },
522            ProviderUnloadAction {
523                kind: UnloadActionKind::ProviderApi,
524                label: "ollama stop fallback".to_string(),
525                command: Some(vec![
526                    "ollama".to_string(),
527                    "stop".to_string(),
528                    profile.model_ref.to_string(),
529                ]),
530                http_method: None,
531                endpoint: None,
532                body_json: None,
533                required: false,
534            },
535        ],
536        UnloadStrategy::VllmSleep => {
537            let endpoint = resolve_provider_endpoint(&profile.unload_strategy, provider_endpoint);
538            vec![ProviderUnloadAction {
539                kind: UnloadActionKind::Sleep,
540                label: "vLLM sleep mode".to_string(),
541                command: None,
542                http_method: Some("POST".to_string()),
543                endpoint: Some(endpoint),
544                body_json: None,
545                required: true,
546            }]
547        }
548        UnloadStrategy::ProcessExit => provider_pid
549            .map(|pid| vec![process_exit_action(pid, "terminate isolated model worker")])
550            .unwrap_or_else(|| {
551                vec![ProviderUnloadAction {
552                    kind: UnloadActionKind::ProcessExit,
553                    label: "terminate isolated model worker".to_string(),
554                    command: None,
555                    http_method: None,
556                    endpoint: None,
557                    body_json: None,
558                    required: true,
559                }]
560            }),
561        UnloadStrategy::None => vec![ProviderUnloadAction {
562            kind: UnloadActionKind::Noop,
563            label: "no GPU unload required".to_string(),
564            command: None,
565            http_method: None,
566            endpoint: None,
567            body_json: None,
568            required: false,
569        }],
570    }
571}
572
573/// Outcome of dispatching a single `ProviderUnloadAction`. The planner owns
574/// execution (#kgunloadpost) so any caller — `tsift kg unload`, a future
575/// lease-drop hook, or the lifecycle swap path — gets the same POST behavior
576/// without re-implementing the HTTP fallback chain.
577#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
578pub struct UnloadActionResult {
579    pub label: String,
580    pub executed: bool,
581    pub outcome: String,
582}
583
584/// Pure projection of a `ProviderUnloadAction` into the request that will be
585/// sent. Separated from `execute_unload_request` so the model-tag override
586/// and body-rewrite logic is fully testable without a live HTTP server.
587///
588/// Returns `None` for non-API actions (noop, process-exit, sleep-only) — those
589/// don't carry an HTTP request and are reported as skipped by the dispatcher.
590#[derive(Debug, Clone, PartialEq, Eq)]
591pub struct PreparedUnloadRequest {
592    pub label: String,
593    pub url: String,
594    pub body: String,
595    pub fallback_command: Option<Vec<String>>,
596}
597
598/// Rewrite the model field of an unload action body to honor an explicit
599/// `--model` override. `build_unload_actions` formats the body with the
600/// profile's `model_ref`, so without this rewrite the override would be
601/// silently ignored. Falls back to the original body if JSON parsing fails
602/// (defensive — the planner's body templates are always valid JSON).
603pub fn rewrite_unload_body_model(body_json: &str, resolved_model_tag: &str) -> String {
604    // Skip re-serialization entirely when there is nothing to rewrite so we
605    // preserve byte-for-byte input on the no-op path (avoids serde_json's
606    // alphabetical key reorder and keeps the empty-override case a true passthrough).
607    if resolved_model_tag.is_empty() {
608        return body_json.to_string();
609    }
610    match serde_json::from_str::<serde_json::Value>(body_json) {
611        Ok(mut value) => {
612            if let Some(obj) = value.as_object_mut() {
613                obj.insert(
614                    "model".to_string(),
615                    serde_json::Value::String(resolved_model_tag.to_string()),
616                );
617            }
618            serde_json::to_string(&value).unwrap_or_else(|_| body_json.to_string())
619        }
620        Err(_) => body_json.to_string(),
621    }
622}
623
624/// Project a planned action into a concrete HTTP request, rewriting the body
625/// with the resolved model tag. Returns `None` for non-API actions.
626pub fn prepare_unload_request(
627    action: &ProviderUnloadAction,
628    resolved_model_tag: &str,
629) -> Option<PreparedUnloadRequest> {
630    if action.kind != UnloadActionKind::ProviderApi {
631        return None;
632    }
633    let endpoint = action.endpoint.clone().unwrap_or_default();
634    let url = normalize_unload_url(&endpoint);
635    let body = match action.body_json.as_deref() {
636        Some(template) => rewrite_unload_body_model(template, resolved_model_tag),
637        None => format!(
638            r#"{{"model":"{resolved_model_tag}","prompt":"","keep_alive":0}}"#
639        ),
640    };
641    Some(PreparedUnloadRequest {
642        label: action.label.clone(),
643        url,
644        body,
645        fallback_command: action.command.clone(),
646    })
647}
648
649/// Normalize an unload endpoint into the canonical `/api/generate` path used
650/// by Ollama's `keep_alive:0` contract. Tolerates callers that supply the
651/// base host (`http://host:11434`) or the full generate URL.
652pub fn normalize_unload_url(endpoint: &str) -> String {
653    let trimmed = endpoint.trim_end_matches('/');
654    if trimmed.ends_with("/api/generate") {
655        return trimmed.to_string();
656    }
657    format!("{}/api/generate", trimmed)
658}
659
660/// Dispatch a single prepared request: POST the body, fall back to the
661/// subprocess command if the HTTP path fails. Impure — exercised end-to-end
662/// via the `tsift kg unload` smoke; unit tests cover `prepare_unload_request`
663/// and `rewrite_unload_body_model` instead.
664pub fn execute_unload_request(req: &PreparedUnloadRequest) -> UnloadActionResult {
665    match post_unload_http(&req.url, &req.body) {
666        Ok(status) => UnloadActionResult {
667            label: req.label.clone(),
668            executed: true,
669            outcome: format!("HTTP {status}"),
670        },
671        Err(err) => {
672            if let Some(cmd) = &req.fallback_command {
673                let _ = std::process::Command::new(&cmd[0])
674                    .args(&cmd[1..])
675                    .status()
676                    .map_err(|e| {
677                        eprintln!("tsift-local-model: unload fallback command failed: {e}");
678                    });
679                return UnloadActionResult {
680                    label: req.label.clone(),
681                    executed: true,
682                    outcome: format!("POST failed ({err}); ran fallback `{:?}`", cmd),
683                };
684            }
685            UnloadActionResult {
686                label: req.label.clone(),
687                executed: false,
688                outcome: format!("POST failed ({err}); no fallback"),
689            }
690        }
691    }
692}
693
694/// Fire-and-forget unload for callers that don't have a full action plan —
695/// used by `tsift kg smoke --unload` and any future lease-drop hook that just
696/// needs to push a single model out of VRAM. Builds a minimal Ollama
697/// `keep_alive:0` request and dispatches it.
698pub fn unload_model_at(host: &str, model_tag: &str) -> UnloadActionResult {
699    let req = PreparedUnloadRequest {
700        label: format!("ollama keep_alive zero for {model_tag}"),
701        url: normalize_unload_url(host),
702        body: format!(r#"{{"model":"{model_tag}","prompt":"","keep_alive":0}}"#),
703        fallback_command: Some(vec![
704            "ollama".to_string(),
705            "stop".to_string(),
706            model_tag.to_string(),
707        ]),
708    };
709    execute_unload_request(&req)
710}
711
712/// Plan + execute a full unload action chain in one call. The default path
713/// for callers that don't need to inspect the prepared request. Skips non-API
714/// actions (noop, process-exit) by reporting them as not executed.
715///
716/// Chain semantics: once a `required` action succeeds, subsequent actions are
717/// skipped (they are fallbacks that exist only for the case where the primary
718/// unload path failed). This prevents the redundant `ollama stop` fallback
719/// from running after a successful `keep_alive:0` POST.
720pub fn execute_unload_actions(
721    actions: &[ProviderUnloadAction],
722    resolved_model_tag: &str,
723) -> Vec<UnloadActionResult> {
724    let mut results = Vec::with_capacity(actions.len());
725    let mut required_succeeded = false;
726    for action in actions {
727        if required_succeeded {
728            results.push(UnloadActionResult {
729                label: action.label.clone(),
730                executed: false,
731                outcome: "skipped: prior required unload succeeded".to_string(),
732            });
733            continue;
734        }
735        let result = match prepare_unload_request(action, resolved_model_tag) {
736            Some(req) => execute_unload_request(&req),
737            None => UnloadActionResult {
738                label: action.label.clone(),
739                executed: false,
740                outcome: "skipped: non-API action".to_string(),
741            },
742        };
743        if result.executed && action.required {
744            required_succeeded = true;
745        }
746        results.push(result);
747    }
748    results
749}
750
751fn post_unload_http(url: &str, body: &str) -> anyhow::Result<String> {
752    use std::time::Duration;
753
754    let payload = serde_json::from_str::<serde_json::Value>(body)
755        .unwrap_or(serde_json::Value::Null);
756    let agent = ureq::Agent::config_builder()
757        .http_status_as_error(false)
758        .timeout_global(Some(Duration::from_secs(30)))
759        .build()
760        .new_agent();
761    let mut response = agent
762        .post(url)
763        .send_json(payload)
764        .with_context(|| format!("posting unload to {url}"))?;
765    let status = response.status();
766    let text = response
767        .body_mut()
768        .read_to_string()
769        .with_context(|| format!("reading unload response (HTTP {status})"))?;
770    if !status.is_success() {
771        bail!("unload HTTP {status}: {}", truncate_str_local(&text, 200));
772    }
773    Ok(format!("{status}"))
774}
775
776fn truncate_str_local(s: &str, max: usize) -> String {
777    if s.len() <= max {
778        s.to_string()
779    } else {
780        format!("{}…", &s[..max])
781    }
782}
783
784/// Resolve a provider endpoint for a given unload strategy.
785///
786/// Precedence (highest first): explicit `--provider-endpoint` value →
787/// strategy-specific env var (`TSIFT_LLAMA_CPP_ENDPOINT` /
788/// `TSIFT_OLLAMA_ENDPOINT` / `TSIFT_VLLM_ENDPOINT`) → compile-time default.
789///
790/// Returns an empty string for strategies that do not use an HTTP endpoint
791/// (`ProcessExit`, `None`); callers should not consult the value in those arms.
792pub fn resolve_provider_endpoint(strategy: &UnloadStrategy, explicit: Option<&str>) -> String {
793    if let Some(explicit) = explicit
794        && !explicit.trim().is_empty()
795    {
796        return explicit.to_string();
797    }
798    let (env_var, default): (&str, &str) = match strategy {
799        UnloadStrategy::LlamaCppRouterUnload => {
800            (LLAMA_CPP_ENDPOINT_ENV_VAR, DEFAULT_LLAMA_CPP_ENDPOINT)
801        }
802        UnloadStrategy::OllamaKeepAliveZero => (OLLAMA_ENDPOINT_ENV_VAR, DEFAULT_OLLAMA_ENDPOINT),
803        UnloadStrategy::VllmSleep => (VLLM_ENDPOINT_ENV_VAR, DEFAULT_VLLM_ENDPOINT),
804        UnloadStrategy::ProcessExit | UnloadStrategy::None => return String::new(),
805    };
806    if let Ok(value) = std::env::var(env_var)
807        && !value.trim().is_empty()
808    {
809        return value;
810    }
811    default.to_string()
812}
813
814pub fn build_lifecycle_report(
815    profile: ModelProfile,
816    pre_load_gpu_probe: GpuProbe,
817    post_unload_gpu_probe: GpuProbe,
818    provider_endpoint: Option<String>,
819    provider_pid: Option<u32>,
820    idle_ttl_seconds: u64,
821    tolerance_mib: u64,
822) -> LocalModelLifecycleReport {
823    let lease = build_local_model_lease(
824        profile,
825        pre_load_gpu_probe.clone(),
826        provider_endpoint,
827        provider_pid,
828        idle_ttl_seconds,
829    );
830    let cleanup = evaluate_vram_cleanup(&pre_load_gpu_probe, &post_unload_gpu_probe, tolerance_mib);
831    let mut notes = vec![match lease.mode {
832        LeaseMode::Exclusive => {
833            "large extractor profile requires an exclusive local-model lease".to_string()
834        }
835        LeaseMode::Shared => "small model profile can share GPU when the margin fits".to_string(),
836        LeaseMode::CpuOrHash => "profile does not require GPU VRAM".to_string(),
837    }];
838    if !cleanup.cleanup_proven {
839        notes.push(
840            "future KG runs should fail if cleanup remains unproven after required unload actions"
841                .to_string(),
842        );
843    }
844
845    LocalModelLifecycleReport {
846        lease,
847        post_unload_gpu_probe,
848        cleanup,
849        notes,
850    }
851}
852
853pub fn evaluate_vram_cleanup(
854    pre_load_gpu_probe: &GpuProbe,
855    post_unload_gpu_probe: &GpuProbe,
856    tolerance_mib: u64,
857) -> VramCleanupEvaluation {
858    let pre_used_mib = pre_load_gpu_probe.used_vram_mib;
859    let post_used_mib = post_unload_gpu_probe.used_vram_mib;
860    let allowed_post_used_mib = pre_used_mib.map(|used| used.saturating_add(tolerance_mib));
861    let used_delta_mib = match (pre_used_mib, post_used_mib) {
862        (Some(pre), Some(post)) => Some(post as i64 - pre as i64),
863        _ => None,
864    };
865
866    if !pre_load_gpu_probe.available
867        || !post_unload_gpu_probe.available
868        || pre_used_mib.is_none()
869        || post_used_mib.is_none()
870    {
871        return VramCleanupEvaluation {
872            status: VramCleanupStatus::ProbeUnavailable,
873            cleanup_proven: false,
874            pre_used_mib,
875            post_used_mib,
876            allowed_post_used_mib,
877            used_delta_mib,
878            external_process_delta_mib: 0,
879            blocking_processes: Vec::new(),
880            reason: "pre-load or post-unload GPU probe is unavailable".to_string(),
881        };
882    }
883
884    let pre_used = pre_used_mib.unwrap();
885    let post_used = post_used_mib.unwrap();
886    let allowed = allowed_post_used_mib.unwrap();
887    if post_used <= allowed {
888        return VramCleanupEvaluation {
889            status: VramCleanupStatus::Proven,
890            cleanup_proven: true,
891            pre_used_mib,
892            post_used_mib,
893            allowed_post_used_mib,
894            used_delta_mib,
895            external_process_delta_mib: 0,
896            blocking_processes: Vec::new(),
897            reason: format!(
898                "post-unload VRAM {post_used} MiB is within {tolerance_mib} MiB of baseline {pre_used} MiB"
899            ),
900        };
901    }
902
903    let blocking_processes = post_unload_gpu_probe
904        .processes
905        .iter()
906        .filter(|process| is_tsift_model_process(process))
907        .cloned()
908        .collect::<Vec<_>>();
909    let external_process_delta_mib =
910        external_process_delta_mib(pre_load_gpu_probe, post_unload_gpu_probe);
911
912    if blocking_processes.is_empty()
913        && post_used <= allowed.saturating_add(external_process_delta_mib)
914    {
915        return VramCleanupEvaluation {
916            status: VramCleanupStatus::ProvenByExternalAccounting,
917            cleanup_proven: true,
918            pre_used_mib,
919            post_used_mib,
920            allowed_post_used_mib,
921            used_delta_mib,
922            external_process_delta_mib,
923            blocking_processes,
924            reason: format!(
925                "post-unload VRAM increase is accounted for by {external_process_delta_mib} MiB of non-tsift GPU processes"
926            ),
927        };
928    }
929
930    VramCleanupEvaluation {
931        status: VramCleanupStatus::NotProven,
932        cleanup_proven: false,
933        pre_used_mib,
934        post_used_mib,
935        allowed_post_used_mib,
936        used_delta_mib,
937        external_process_delta_mib,
938        blocking_processes,
939        reason: format!(
940            "post-unload VRAM {post_used} MiB exceeds allowed {allowed} MiB and cleanup is not externally accounted for"
941        ),
942    }
943}
944
945pub fn rank_profiles_for_role(
946    profiles: &[ModelProfile],
947    probe: &GpuProbe,
948    role: ModelRole,
949) -> Vec<ProfileSelection> {
950    profiles
951        .iter()
952        .filter(|profile| profile.supports_role(&role))
953        .map(|profile| selection_for_profile(profile, probe))
954        .collect()
955}
956
957pub fn format_status_human(report: &LocalModelStatusReport) -> String {
958    let mut out = String::new();
959    out.push_str("Local model status\n");
960    if report.gpu_probe.available {
961        out.push_str(&format!(
962            "GPU: {} | VRAM: {} MiB used / {} MiB total ({} MiB free)\n",
963            report.gpu_probe.gpu_name.as_deref().unwrap_or("unknown"),
964            format_optional_u64(report.gpu_probe.used_vram_mib),
965            format_optional_u64(report.gpu_probe.total_vram_mib),
966            format_optional_u64(report.gpu_probe.free_vram_mib)
967        ));
968    } else {
969        out.push_str(&format!(
970            "GPU: unavailable ({})\n",
971            report.gpu_probe.error.as_deref().unwrap_or("unknown error")
972        ));
973    }
974    out.push_str(&format!(
975        "Recommended extractor: {}\n",
976        report
977            .recommended_extractor
978            .as_deref()
979            .unwrap_or("none selectable")
980    ));
981    out.push_str(&format!(
982        "Recommended embedding: {}\n",
983        report
984            .recommended_embedding
985            .as_deref()
986            .unwrap_or("none selectable")
987    ));
988    out.push_str("\nExtractor profiles:\n");
989    for selection in &report.extractor_profiles {
990        out.push_str(&format!(
991            "- {} [{} MiB est]: {} ({})\n",
992            selection.profile.id,
993            selection.profile.estimated_total_mib(),
994            if selection.selectable {
995                "selectable"
996            } else {
997                "blocked"
998            },
999            selection.reason
1000        ));
1001    }
1002    out.push_str("\nEmbedding profiles:\n");
1003    for selection in &report.embedding_profiles {
1004        out.push_str(&format!(
1005            "- {} [{} MiB est]: {} ({})\n",
1006            selection.profile.id,
1007            selection.profile.estimated_total_mib(),
1008            if selection.selectable {
1009                "selectable"
1010            } else {
1011                "blocked"
1012            },
1013            selection.reason
1014        ));
1015    }
1016    out
1017}
1018
1019pub fn format_lifecycle_human(report: &LocalModelLifecycleReport) -> String {
1020    let mut out = String::new();
1021    out.push_str("Local model lifecycle\n");
1022    out.push_str(&format!(
1023        "Profile: {} ({})\n",
1024        report.lease.profile.id, report.lease.profile.label
1025    ));
1026    out.push_str(&format!(
1027        "Lease: {} | mode: {:?} | idle TTL: {}s\n",
1028        report.lease.lease_id, report.lease.mode, report.lease.idle_ttl_seconds
1029    ));
1030    out.push_str(&format!(
1031        "Pre-load VRAM: {} MiB used\n",
1032        format_optional_u64(report.lease.pre_load_gpu_probe.used_vram_mib)
1033    ));
1034    out.push_str(&format!(
1035        "Post-unload VRAM: {} MiB used\n",
1036        format_optional_u64(report.post_unload_gpu_probe.used_vram_mib)
1037    ));
1038    out.push_str(&format!(
1039        "Cleanup: {:?} ({})\n",
1040        report.cleanup.status, report.cleanup.reason
1041    ));
1042    out.push_str("\nRequired unload actions:\n");
1043    for action in &report.lease.unload_actions {
1044        out.push_str(&format!(
1045            "- {}: {}{}\n",
1046            if action.required {
1047                "required"
1048            } else {
1049                "fallback"
1050            },
1051            action.label,
1052            format_action_detail(action)
1053        ));
1054    }
1055    if !report.cleanup.blocking_processes.is_empty() {
1056        out.push_str("\nBlocking GPU processes:\n");
1057        for process in &report.cleanup.blocking_processes {
1058            out.push_str(&format!(
1059                "- pid={} name={} used={} MiB\n",
1060                process
1061                    .pid
1062                    .map(|pid| pid.to_string())
1063                    .unwrap_or_else(|| "unknown".to_string()),
1064                process.process_name,
1065                format_optional_u64(process.used_memory_mib)
1066            ));
1067        }
1068    }
1069    out
1070}
1071
1072fn format_action_detail(action: &ProviderUnloadAction) -> String {
1073    if let Some(command) = &action.command {
1074        return format!(" | command: {}", command.join(" "));
1075    }
1076    if let Some(endpoint) = &action.endpoint {
1077        return format!(
1078            " | {} {}{}",
1079            action.http_method.as_deref().unwrap_or("POST"),
1080            endpoint,
1081            action
1082                .body_json
1083                .as_ref()
1084                .map(|body| format!(" body={body}"))
1085                .unwrap_or_default()
1086        );
1087    }
1088    String::new()
1089}
1090
1091fn selection_for_profile(profile: &ModelProfile, probe: &GpuProbe) -> ProfileSelection {
1092    if profile.concurrency == ConcurrencyClass::CpuOrHash {
1093        return ProfileSelection {
1094            profile: profile.clone(),
1095            selectable: true,
1096            reason: "does not require GPU VRAM".to_string(),
1097        };
1098    }
1099
1100    let Some(free_vram_mib) = probe.free_vram_mib else {
1101        return ProfileSelection {
1102            profile: profile.clone(),
1103            selectable: false,
1104            reason: "free VRAM unknown".to_string(),
1105        };
1106    };
1107
1108    let required = profile.estimated_total_mib();
1109    if required <= free_vram_mib {
1110        ProfileSelection {
1111            profile: profile.clone(),
1112            selectable: true,
1113            reason: format!("estimated {required} MiB fits in {free_vram_mib} MiB free"),
1114        }
1115    } else {
1116        ProfileSelection {
1117            profile: profile.clone(),
1118            selectable: false,
1119            reason: format!("estimated {required} MiB exceeds {free_vram_mib} MiB free"),
1120        }
1121    }
1122}
1123
1124fn process_exit_action(pid: u32, label: &str) -> ProviderUnloadAction {
1125    ProviderUnloadAction {
1126        kind: UnloadActionKind::ProcessExit,
1127        label: label.to_string(),
1128        command: Some(vec![
1129            "kill".to_string(),
1130            "-TERM".to_string(),
1131            pid.to_string(),
1132        ]),
1133        http_method: None,
1134        endpoint: None,
1135        body_json: None,
1136        required: false,
1137    }
1138}
1139
1140fn external_process_delta_mib(
1141    pre_load_gpu_probe: &GpuProbe,
1142    post_unload_gpu_probe: &GpuProbe,
1143) -> u64 {
1144    post_unload_gpu_probe
1145        .processes
1146        .iter()
1147        .filter(|process| !is_tsift_model_process(process))
1148        .map(|process| {
1149            let before = matching_pre_process(pre_load_gpu_probe, process)
1150                .and_then(|pre| pre.used_memory_mib)
1151                .unwrap_or(0);
1152            process.used_memory_mib.unwrap_or(0).saturating_sub(before)
1153        })
1154        .sum()
1155}
1156
1157fn matching_pre_process<'a>(
1158    pre_load_gpu_probe: &'a GpuProbe,
1159    post_process: &GpuProcess,
1160) -> Option<&'a GpuProcess> {
1161    if let Some(pid) = post_process.pid
1162        && let Some(process) = pre_load_gpu_probe
1163            .processes
1164            .iter()
1165            .find(|candidate| candidate.pid == Some(pid))
1166    {
1167        return Some(process);
1168    }
1169    pre_load_gpu_probe
1170        .processes
1171        .iter()
1172        .find(|candidate| candidate.process_name == post_process.process_name)
1173}
1174
1175fn is_tsift_model_process(process: &GpuProcess) -> bool {
1176    let name = process.process_name.to_ascii_lowercase();
1177    name.contains("tsift")
1178        || name.contains("llama")
1179        || name.contains("ollama")
1180        || name.contains("vllm")
1181        || name.contains("ggml")
1182}
1183
1184pub fn current_unix_seconds() -> u64 {
1185    SystemTime::now()
1186        .duration_since(UNIX_EPOCH)
1187        .map(|duration| duration.as_secs())
1188        .unwrap_or(0)
1189}
1190
1191fn parse_gpu_query(line: &str) -> Result<GpuProbe> {
1192    let parts = line.split(',').map(str::trim).collect::<Vec<_>>();
1193    if parts.len() != 5 {
1194        anyhow::bail!("unexpected nvidia-smi gpu query row: {line}");
1195    }
1196    Ok(GpuProbe {
1197        timestamp_unix_seconds: Some(current_unix_seconds()),
1198        available: true,
1199        gpu_name: Some(parts[0].to_string()),
1200        driver_version: Some(parts[1].to_string()),
1201        total_vram_mib: Some(parse_optional_u64(parts[2]).context("parse total VRAM")?),
1202        used_vram_mib: Some(parse_optional_u64(parts[3]).context("parse used VRAM")?),
1203        free_vram_mib: Some(parse_optional_u64(parts[4]).context("parse free VRAM")?),
1204        processes: Vec::new(),
1205        error: None,
1206    })
1207}
1208
1209fn query_nvidia_compute_processes() -> Vec<GpuProcess> {
1210    let Ok(output) = Command::new("nvidia-smi")
1211        .args([
1212            "--query-compute-apps=pid,process_name,used_memory",
1213            "--format=csv,noheader,nounits",
1214        ])
1215        .output()
1216    else {
1217        return Vec::new();
1218    };
1219    if !output.status.success() {
1220        return Vec::new();
1221    }
1222    String::from_utf8_lossy(&output.stdout)
1223        .lines()
1224        .filter_map(parse_process_query)
1225        .collect()
1226}
1227
1228fn parse_process_query(line: &str) -> Option<GpuProcess> {
1229    let parts = line.split(',').map(str::trim).collect::<Vec<_>>();
1230    if parts.len() != 3 || parts.iter().all(|part| part.is_empty()) {
1231        return None;
1232    }
1233    Some(GpuProcess {
1234        pid: parts[0].parse::<u32>().ok(),
1235        process_name: parts[1].to_string(),
1236        used_memory_mib: parse_optional_u64(parts[2]).ok(),
1237    })
1238}
1239
1240fn parse_optional_u64(input: &str) -> Result<u64> {
1241    let cleaned = input.trim().trim_end_matches("MiB").trim();
1242    cleaned
1243        .parse::<u64>()
1244        .with_context(|| format!("parse integer from {input:?}"))
1245}
1246
1247fn format_optional_u64(value: Option<u64>) -> String {
1248    value
1249        .map(|value| value.to_string())
1250        .unwrap_or_else(|| "unknown".to_string())
1251}
1252
1253// ============================================================================
1254// Cooperative GPU lease registry (#gctrl1)
1255//
1256// A file-backed registry of who currently holds a GPU-bound local model
1257// profile. Cooperative (no daemon): producers check the file before probing
1258// the GPU, prune stale leases (dead pid or past idle TTL), and either acquire
1259// the slot or report a conflict with the live holder.
1260//
1261// The registry is keyed by `profile_id` and holds a list of `GpuLeaseRecord`
1262// holders. `Exclusive` profiles allow at most one live holder; `Shared`
1263// profiles allow many; `CpuOrHash` profiles bypass the registry entirely
1264// because they do not consume GPU VRAM.
1265// ============================================================================
1266
1267/// Cooperative GPU lease registry file format version.
1268pub const LEASE_REGISTRY_VERSION: u32 = 1;
1269/// Default idle TTL (0 = no TTL-based staleness, only pid-dead pruning).
1270pub const DEFAULT_LEASE_TTL_SECONDS: u64 = 0;
1271/// Environment variable override for the lease registry file path.
1272pub const LEASE_FILE_ENV_VAR: &str = "TSIFT_LEASE_FILE";
1273
1274/// One held lease on a profile, written to the cooperative registry file.
1275#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1276pub struct GpuLeaseRecord {
1277    pub profile_id: String,
1278    pub holder_pid: u32,
1279    pub holder_command: String,
1280    /// Acquire time, which also serves as the last-heartbeat timestamp: a
1281    /// re-acquire by the same pid or an explicit `renew` slides it forward, so
1282    /// `idle_ttl_seconds`-based staleness is measured against the most recent
1283    /// heartbeat rather than the original acquire.
1284    pub acquired_at_unix_seconds: u64,
1285    pub lease_mode: LeaseMode,
1286    pub vram_baseline_mib: u64,
1287    pub idle_ttl_seconds: u64,
1288    pub notes: Vec<String>,
1289}
1290
1291/// File-backed cooperative registry: `{ version, leases: { profile_id: [record, ...] } }`.
1292#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1293pub struct GpuLeaseRegistry {
1294    pub version: u32,
1295    pub leases: BTreeMap<String, Vec<GpuLeaseRecord>>,
1296}
1297
1298impl Default for GpuLeaseRegistry {
1299    fn default() -> Self {
1300        Self {
1301            version: LEASE_REGISTRY_VERSION,
1302            leases: BTreeMap::new(),
1303        }
1304    }
1305}
1306
1307#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1308pub enum GpuLeaseAcquisitionStatus {
1309    /// Fresh acquire on a free slot.
1310    Acquired,
1311    /// Same holder pid already held the slot; timestamp/baseline refreshed.
1312    Refreshed,
1313    /// Previous holder was stale (pid gone or TTL expired); slot reclaimed.
1314    ReclaimedStale,
1315    /// Profile is `CpuOrHash`; no registry entry needed.
1316    CpuOrHashBypass,
1317    /// Another live holder owns the slot.
1318    Conflict,
1319}
1320
1321#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1322pub struct GpuLeaseConflict {
1323    pub profile_id: String,
1324    pub holder_pid: u32,
1325    pub holder_command: String,
1326    pub acquired_at_unix_seconds: u64,
1327    pub lease_mode: LeaseMode,
1328}
1329
1330#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1331pub struct GpuLeaseAcquisition {
1332    pub profile_id: String,
1333    pub holder_pid: u32,
1334    pub status: GpuLeaseAcquisitionStatus,
1335    pub record: Option<GpuLeaseRecord>,
1336    pub conflict: Option<GpuLeaseConflict>,
1337    /// Stale records pruned during this acquire (cleared from the registry).
1338    pub reclaimed: Vec<GpuLeaseRecord>,
1339}
1340
1341#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1342pub enum GpuLeaseReleaseOutcome {
1343    /// This holder's lease was removed.
1344    Released,
1345    /// Profile exists but this pid was not among its holders.
1346    NotHeld,
1347    /// No entry for the profile at all.
1348    ProfileAbsent,
1349}
1350
1351#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1352pub struct GpuLeaseRelease {
1353    pub profile_id: String,
1354    pub holder_pid: u32,
1355    pub outcome: GpuLeaseReleaseOutcome,
1356    /// Number of remaining live holders for the profile after release.
1357    pub remaining_holders: u32,
1358}
1359
1360/// Resolve the cooperative lease registry file path.
1361///
1362/// Order: explicit `override_path` → `$TSIFT_LEASE_FILE` →
1363/// `$XDG_STATE_HOME/tsift/gpu-lease.json` → `~/.tsift/gpu-lease.json` →
1364/// `./.tsift/gpu-lease.json` if no home directory can be resolved.
1365pub fn resolve_lease_file(override_path: Option<&Path>) -> PathBuf {
1366    if let Some(path) = override_path {
1367        return path.to_path_buf();
1368    }
1369    if let Ok(path) = std::env::var(LEASE_FILE_ENV_VAR) {
1370        return PathBuf::from(path);
1371    }
1372    if let Ok(state_dir) = std::env::var("XDG_STATE_HOME")
1373        && !state_dir.is_empty()
1374    {
1375        return PathBuf::from(state_dir)
1376            .join("tsift")
1377            .join("gpu-lease.json");
1378    }
1379    if let Ok(home) = std::env::var("HOME")
1380        && !home.is_empty()
1381    {
1382        return PathBuf::from(home).join(".tsift").join("gpu-lease.json");
1383    }
1384    PathBuf::from("./.tsift/gpu-lease.json")
1385}
1386
1387/// Best-effort pid-liveness check via `kill -0`.
1388///
1389/// A pid equal to the current process is always considered alive. Pid 0 is
1390/// treated as missing/unknown and reported as not alive so callers can use 0
1391/// as a sentinel for "no pid recorded".
1392pub fn is_pid_alive(pid: u32) -> bool {
1393    if pid == 0 {
1394        return false;
1395    }
1396    if pid == std::process::id() {
1397        return true;
1398    }
1399    match Command::new("kill").arg("-0").arg(pid.to_string()).output() {
1400        Ok(output) => output.status.success(),
1401        Err(_) => false,
1402    }
1403}
1404
1405/// Sidecar advisory-lock path for a registry file (`<registry>.lock`).
1406///
1407/// A dedicated lock file (rather than locking the registry itself) keeps the
1408/// lock independent of the atomic temp-file + rename write, which replaces the
1409/// registry inode on every write.
1410pub fn registry_lock_path(path: &Path) -> PathBuf {
1411    let mut name = path.as_os_str().to_os_string();
1412    name.push(".lock");
1413    PathBuf::from(name)
1414}
1415
1416/// Run `op` while holding an exclusive advisory lock on the registry's sidecar
1417/// lock file.
1418///
1419/// The cooperative registry is mutated with a read → apply → write cycle. The
1420/// atomic temp-file + rename in [`write_lease_registry`] makes each *write*
1421/// atomic, but two processes can still interleave read/apply/write and lose an
1422/// update (TOCTOU). Holding an OS advisory lock across the whole cycle
1423/// serializes concurrent acquire/release/renew/reap across processes, and the
1424/// kernel releases the lock automatically if the holder dies mid-cycle — so a
1425/// crashed holder can never wedge the registry.
1426fn with_registry_lock<T>(path: &Path, op: impl FnOnce() -> Result<T>) -> Result<T> {
1427    use fs4::fs_std::FileExt;
1428    let lock_path = registry_lock_path(path);
1429    if let Some(parent) = lock_path.parent()
1430        && !parent.as_os_str().is_empty()
1431    {
1432        fs::create_dir_all(parent).context("create lease registry lock parent")?;
1433    }
1434    let lock_file = fs::OpenOptions::new()
1435        .create(true)
1436        .read(true)
1437        .write(true)
1438        .truncate(false)
1439        .open(&lock_path)
1440        .with_context(|| format!("open lease registry lock {}", lock_path.display()))?;
1441    lock_file
1442        .lock_exclusive()
1443        .context("acquire exclusive lease registry lock")?;
1444    let result = op();
1445    // Best-effort unlock; the lock is also released when `lock_file` drops or
1446    // the process exits.
1447    let _ = FileExt::unlock(&lock_file);
1448    result
1449}
1450
1451/// Read the lease registry, returning an empty default when the file is missing.
1452pub fn read_lease_registry(path: &Path) -> Result<GpuLeaseRegistry> {
1453    match fs::read_to_string(path) {
1454        Ok(contents) => {
1455            if contents.trim().is_empty() {
1456                return Ok(GpuLeaseRegistry::default());
1457            }
1458            serde_json::from_str(&contents).context("parse gpu lease registry")
1459        }
1460        Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
1461            Ok(GpuLeaseRegistry::default())
1462        }
1463        Err(error) => Err(error).context("read gpu lease registry"),
1464    }
1465}
1466
1467/// Atomically write the lease registry (temp file + rename).
1468pub fn write_lease_registry(path: &Path, registry: &GpuLeaseRegistry) -> Result<()> {
1469    if let Some(parent) = path.parent()
1470        && !parent.as_os_str().is_empty()
1471    {
1472        fs::create_dir_all(parent).context("create lease registry parent")?;
1473    }
1474    let payload = serde_json::to_string_pretty(registry).context("serialize lease registry")?;
1475    let temp_path = path.with_extension(format!(
1476        "json.tmp.{}.{}",
1477        std::process::id(),
1478        current_unix_seconds()
1479    ));
1480    let mut handle = fs::File::create(&temp_path).context("create lease registry temp file")?;
1481    handle
1482        .write_all(payload.as_bytes())
1483        .context("write lease registry temp file")?;
1484    handle.sync_all().context("sync lease registry temp file")?;
1485    drop(handle);
1486    fs::rename(&temp_path, path).context("rename lease registry into place")?;
1487    Ok(())
1488}
1489
1490/// Prune stale holders from the registry in place.
1491///
1492/// A holder is stale when its pid is no longer alive, or when its
1493/// `idle_ttl_seconds > 0` and the lease age exceeds the TTL.
1494///
1495/// Returns the records that were pruned.
1496pub fn prune_stale_leases(
1497    registry: &mut GpuLeaseRegistry,
1498    now: u64,
1499    is_alive: impl Fn(u32) -> bool,
1500) -> Vec<GpuLeaseRecord> {
1501    let mut pruned = Vec::new();
1502    let mut empty_keys = Vec::new();
1503    for (profile_id, holders) in registry.leases.iter_mut() {
1504        let mut kept = Vec::with_capacity(holders.len());
1505        for record in holders.drain(..) {
1506            let pid_dead = !is_alive(record.holder_pid);
1507            let ttl_expired = record.idle_ttl_seconds > 0
1508                && now.saturating_sub(record.acquired_at_unix_seconds) > record.idle_ttl_seconds;
1509            if pid_dead || ttl_expired {
1510                pruned.push(record);
1511            } else {
1512                kept.push(record);
1513            }
1514        }
1515        if kept.is_empty() {
1516            empty_keys.push(profile_id.clone());
1517        }
1518        *holders = kept;
1519    }
1520    for key in empty_keys {
1521        registry.leases.remove(&key);
1522    }
1523    pruned
1524}
1525
1526/// Apply an acquire to the registry in place.
1527///
1528/// Pure logic; the file I/O wrapper is `acquire_lease`. The `is_alive` closure
1529/// lets tests inject a deterministic liveness check.
1530#[allow(clippy::too_many_arguments)]
1531pub fn apply_acquire(
1532    registry: &mut GpuLeaseRegistry,
1533    profile: &ModelProfile,
1534    holder_pid: u32,
1535    holder_command: &str,
1536    vram_baseline_mib: u64,
1537    idle_ttl_seconds: u64,
1538    now: u64,
1539    is_alive: impl Fn(u32) -> bool,
1540) -> GpuLeaseAcquisition {
1541    if profile.concurrency == ConcurrencyClass::CpuOrHash {
1542        return GpuLeaseAcquisition {
1543            profile_id: profile.id.to_string(),
1544            holder_pid,
1545            status: GpuLeaseAcquisitionStatus::CpuOrHashBypass,
1546            record: None,
1547            conflict: None,
1548            reclaimed: Vec::new(),
1549        };
1550    }
1551
1552    let reclaimed = prune_stale_leases(registry, now, &is_alive);
1553    let mode = lease_mode_for_profile(profile);
1554    let entry = registry.leases.entry(profile.id.to_string()).or_default();
1555    let already_held = entry
1556        .iter()
1557        .position(|record| record.holder_pid == holder_pid);
1558
1559    let record = GpuLeaseRecord {
1560        profile_id: profile.id.to_string(),
1561        holder_pid,
1562        holder_command: holder_command.to_string(),
1563        acquired_at_unix_seconds: now,
1564        lease_mode: mode.clone(),
1565        vram_baseline_mib,
1566        idle_ttl_seconds,
1567        notes: Vec::new(),
1568    };
1569
1570    let status = if let Some(index) = already_held {
1571        entry[index] = record.clone();
1572        GpuLeaseAcquisitionStatus::Refreshed
1573    } else {
1574        match mode {
1575            LeaseMode::Exclusive => {
1576                if let Some(blocker) = entry.first() {
1577                    return GpuLeaseAcquisition {
1578                        profile_id: profile.id.to_string(),
1579                        holder_pid,
1580                        status: GpuLeaseAcquisitionStatus::Conflict,
1581                        record: None,
1582                        conflict: Some(GpuLeaseConflict {
1583                            profile_id: profile.id.to_string(),
1584                            holder_pid: blocker.holder_pid,
1585                            holder_command: blocker.holder_command.clone(),
1586                            acquired_at_unix_seconds: blocker.acquired_at_unix_seconds,
1587                            lease_mode: blocker.lease_mode.clone(),
1588                        }),
1589                        reclaimed,
1590                    };
1591                }
1592                entry.push(record.clone());
1593                if reclaimed
1594                    .iter()
1595                    .any(|pruned| pruned.profile_id == profile.id)
1596                {
1597                    GpuLeaseAcquisitionStatus::ReclaimedStale
1598                } else {
1599                    GpuLeaseAcquisitionStatus::Acquired
1600                }
1601            }
1602            LeaseMode::Shared => {
1603                entry.push(record.clone());
1604                if reclaimed
1605                    .iter()
1606                    .any(|pruned| pruned.profile_id == profile.id)
1607                {
1608                    GpuLeaseAcquisitionStatus::ReclaimedStale
1609                } else {
1610                    GpuLeaseAcquisitionStatus::Acquired
1611                }
1612            }
1613            LeaseMode::CpuOrHash => GpuLeaseAcquisitionStatus::CpuOrHashBypass,
1614        }
1615    };
1616
1617    GpuLeaseAcquisition {
1618        profile_id: profile.id.to_string(),
1619        holder_pid,
1620        status,
1621        record: Some(record),
1622        conflict: None,
1623        reclaimed,
1624    }
1625}
1626
1627/// Apply a release to the registry in place.
1628pub fn apply_release(
1629    registry: &mut GpuLeaseRegistry,
1630    profile_id: &str,
1631    holder_pid: u32,
1632    now: u64,
1633    is_alive: impl Fn(u32) -> bool,
1634) -> GpuLeaseRelease {
1635    let _ = prune_stale_leases(registry, now, &is_alive);
1636    let Some(holders) = registry.leases.get_mut(profile_id) else {
1637        return GpuLeaseRelease {
1638            profile_id: profile_id.to_string(),
1639            holder_pid,
1640            outcome: GpuLeaseReleaseOutcome::ProfileAbsent,
1641            remaining_holders: 0,
1642        };
1643    };
1644    let before = holders.len();
1645    holders.retain(|record| record.holder_pid != holder_pid);
1646    let removed = before - holders.len();
1647    let remaining = holders.len() as u32;
1648    if holders.is_empty() {
1649        // Borrow on `holders` ends here; safe to mutate the map again.
1650        registry.leases.remove(profile_id);
1651    }
1652    let outcome = if removed == 0 {
1653        GpuLeaseReleaseOutcome::NotHeld
1654    } else {
1655        GpuLeaseReleaseOutcome::Released
1656    };
1657    GpuLeaseRelease {
1658        profile_id: profile_id.to_string(),
1659        holder_pid,
1660        outcome,
1661        remaining_holders: remaining,
1662    }
1663}
1664
1665/// High-level acquire: read file, prune stale, apply, write file.
1666pub fn acquire_lease(
1667    profile_id: &str,
1668    holder_pid: u32,
1669    holder_command: &str,
1670    vram_baseline_mib: u64,
1671    idle_ttl_seconds: u64,
1672    now: u64,
1673    path: &Path,
1674) -> Result<GpuLeaseAcquisition> {
1675    let profile = profile_by_id(profile_id)
1676        .with_context(|| format!("unknown local model profile {profile_id:?}"))?;
1677    with_registry_lock(path, || {
1678        let mut registry = read_lease_registry(path)?;
1679        let acquisition = apply_acquire(
1680            &mut registry,
1681            &profile,
1682            holder_pid,
1683            holder_command,
1684            vram_baseline_mib,
1685            idle_ttl_seconds,
1686            now,
1687            is_pid_alive,
1688        );
1689        // CpuOrHash bypass intentionally does not touch the registry file.
1690        if acquisition.status != GpuLeaseAcquisitionStatus::CpuOrHashBypass {
1691            write_lease_registry(path, &registry)?;
1692        }
1693        Ok(acquisition)
1694    })
1695}
1696
1697/// High-level release: read file, prune, drop this holder, write file.
1698pub fn release_lease(
1699    profile_id: &str,
1700    holder_pid: u32,
1701    now: u64,
1702    path: &Path,
1703) -> Result<GpuLeaseRelease> {
1704    with_registry_lock(path, || {
1705        let mut registry = read_lease_registry(path)?;
1706        let release = apply_release(&mut registry, profile_id, holder_pid, now, is_pid_alive);
1707        write_lease_registry(path, &registry)?;
1708        Ok(release)
1709    })
1710}
1711
1712/// Read the registry and return the pruned view. `include_stale` skips the
1713/// pruning pass so the caller can inspect raw state for diagnostics.
1714pub fn show_registry(path: &Path, now: u64, include_stale: bool) -> Result<GpuLeaseRegistry> {
1715    let mut registry = read_lease_registry(path)?;
1716    if !include_stale {
1717        prune_stale_leases(&mut registry, now, is_pid_alive);
1718    }
1719    Ok(registry)
1720}
1721
1722#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1723pub enum GpuLeaseRenewOutcome {
1724    /// The holder's heartbeat timestamp was slid forward to `now`.
1725    Renewed,
1726    /// Profile exists but this pid was not among its holders.
1727    NotHeld,
1728    /// No entry for the profile at all.
1729    ProfileAbsent,
1730}
1731
1732#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1733pub struct GpuLeaseRenew {
1734    pub profile_id: String,
1735    pub holder_pid: u32,
1736    pub outcome: GpuLeaseRenewOutcome,
1737    /// New heartbeat timestamp written for the holder (when renewed).
1738    pub renewed_at_unix_seconds: Option<u64>,
1739}
1740
1741/// Apply a heartbeat renewal in place: slide a live holder's heartbeat
1742/// (`acquired_at_unix_seconds`) forward to `now` so its TTL window restarts.
1743pub fn apply_renew(
1744    registry: &mut GpuLeaseRegistry,
1745    profile_id: &str,
1746    holder_pid: u32,
1747    now: u64,
1748    is_alive: impl Fn(u32) -> bool,
1749) -> GpuLeaseRenew {
1750    let _ = prune_stale_leases(registry, now, &is_alive);
1751    let Some(holders) = registry.leases.get_mut(profile_id) else {
1752        return GpuLeaseRenew {
1753            profile_id: profile_id.to_string(),
1754            holder_pid,
1755            outcome: GpuLeaseRenewOutcome::ProfileAbsent,
1756            renewed_at_unix_seconds: None,
1757        };
1758    };
1759    if let Some(record) = holders
1760        .iter_mut()
1761        .find(|record| record.holder_pid == holder_pid)
1762    {
1763        record.acquired_at_unix_seconds = now;
1764        GpuLeaseRenew {
1765            profile_id: profile_id.to_string(),
1766            holder_pid,
1767            outcome: GpuLeaseRenewOutcome::Renewed,
1768            renewed_at_unix_seconds: Some(now),
1769        }
1770    } else {
1771        GpuLeaseRenew {
1772            profile_id: profile_id.to_string(),
1773            holder_pid,
1774            outcome: GpuLeaseRenewOutcome::NotHeld,
1775            renewed_at_unix_seconds: None,
1776        }
1777    }
1778}
1779
1780/// High-level heartbeat: read file, slide this holder's heartbeat, write file.
1781///
1782/// A long-lived session calls this periodically so its lease is held against
1783/// `idle_ttl_seconds`-based reclamation without re-running probes.
1784pub fn renew_lease(
1785    profile_id: &str,
1786    holder_pid: u32,
1787    now: u64,
1788    path: &Path,
1789) -> Result<GpuLeaseRenew> {
1790    with_registry_lock(path, || {
1791        let mut registry = read_lease_registry(path)?;
1792        let renew = apply_renew(&mut registry, profile_id, holder_pid, now, is_pid_alive);
1793        write_lease_registry(path, &registry)?;
1794        Ok(renew)
1795    })
1796}
1797
1798#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1799pub struct GpuLeaseReap {
1800    /// Stale records pruned this reap (pid dead or TTL expired).
1801    pub reclaimed: Vec<GpuLeaseRecord>,
1802    /// Profile ids whose last live holder was reclaimed this reap — the
1803    /// reference count dropped to zero, so the model can be unloaded.
1804    pub emptied_profiles: Vec<String>,
1805}
1806
1807/// Sweep stale holders (crashed pids or expired TTL) out of the registry and
1808/// report which profiles dropped to zero live holders.
1809///
1810/// This is the crash-reclamation entrypoint: a session that died without
1811/// releasing leaves a pid-dead holder that `reap` clears, and when that was the
1812/// last reference for a profile the profile appears in `emptied_profiles` so
1813/// the caller can unload the now-unreferenced model.
1814pub fn reap_leases(now: u64, path: &Path) -> Result<GpuLeaseReap> {
1815    with_registry_lock(path, || {
1816        let mut registry = read_lease_registry(path)?;
1817        let before: std::collections::BTreeSet<String> =
1818            registry.leases.keys().cloned().collect();
1819        let reclaimed = prune_stale_leases(&mut registry, now, is_pid_alive);
1820        let emptied_profiles: Vec<String> = before
1821            .into_iter()
1822            .filter(|profile_id| !registry.leases.contains_key(profile_id))
1823            .collect();
1824        write_lease_registry(path, &registry)?;
1825        Ok(GpuLeaseReap {
1826            reclaimed,
1827            emptied_profiles,
1828        })
1829    })
1830}
1831
1832/// Human-readable summary of the lease registry.
1833pub fn format_lease_show_human(registry: &GpuLeaseRegistry, now: u64) -> String {
1834    let mut out = String::new();
1835    out.push_str("GPU lease registry\n");
1836    out.push_str(&format!("version: {}\n", registry.version));
1837    if registry.leases.is_empty() {
1838        out.push_str("leases: none\n");
1839        return out;
1840    }
1841    out.push_str(&format!("profiles held: {}\n", registry.leases.len()));
1842    for (profile_id, holders) in &registry.leases {
1843        out.push_str(&format!("\n{profile_id} ({} holder(s)):\n", holders.len()));
1844        for record in holders {
1845            let age = now.saturating_sub(record.acquired_at_unix_seconds);
1846            out.push_str(&format!(
1847                "  pid={} cmd={} mode={:?} acquired={}s ago baseline={} MiB ttl={}s",
1848                record.holder_pid,
1849                record.holder_command,
1850                record.lease_mode,
1851                age,
1852                record.vram_baseline_mib,
1853                record.idle_ttl_seconds
1854            ));
1855            if record.notes.is_empty() {
1856                out.push('\n');
1857            } else {
1858                out.push_str(&format!(" notes={}\n", record.notes.join("; ")));
1859            }
1860        }
1861    }
1862    out
1863}
1864
1865// ============================================================================
1866// Per-call profile preference (#gctrl2)
1867//
1868// Callers (agent-doc cycles, scripts) that want to pin or downgrade the local
1869// model for a single call — without mutating global state — express that as a
1870// `ProfilePreference`. The resolver turns the preference + the live GPU probe
1871// into a concrete `ProfileSelection` plus a `ProfileResolutionSource` saying
1872// how the choice was made. Commands that touch the local model accept the
1873// preference via `--profile`, record it in their response envelope, and will
1874// hand it to the real provider seam once one is wired in.
1875// ============================================================================
1876
1877/// Caller-supplied preference for which local model profile a single call
1878/// should use.
1879#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1880#[serde(tag = "kind", content = "value", rename_all = "snake_case")]
1881pub enum ProfilePreference {
1882    /// No pin — rank by free VRAM (existing behavior).
1883    Auto,
1884    /// Pin to a specific profile id. The resolver still checks VRAM fit and
1885    /// reports `PinnedUnselectable` if the profile would not fit the probe.
1886    Pinned(String),
1887    /// Force the deterministic CPU/hash fallback even if a GPU profile would
1888    /// fit. Use during low-stakes phases of a long agent-doc run.
1889    ForceHash,
1890}
1891
1892impl ProfilePreference {
1893    /// Parse the `--profile <Option<String>>` CLI value.
1894    ///
1895    /// `None` / empty → `Auto`. The literal `"hash"` or the hash profile id
1896    /// (`tsift-local-hash-v1`) → `ForceHash`. Anything else → `Pinned(id)`.
1897    pub fn from_cli(value: Option<&str>) -> Self {
1898        match value.map(str::trim) {
1899            None | Some("") => ProfilePreference::Auto,
1900            Some("hash") | Some("tsift-local-hash-v1") => ProfilePreference::ForceHash,
1901            Some(other) => ProfilePreference::Pinned(other.to_string()),
1902        }
1903    }
1904
1905    pub fn describe(&self) -> String {
1906        match self {
1907            ProfilePreference::Auto => "auto".to_string(),
1908            ProfilePreference::Pinned(id) => format!("pinned:{id}"),
1909            ProfilePreference::ForceHash => "force-hash".to_string(),
1910        }
1911    }
1912}
1913
1914/// How a resolved profile was chosen.
1915#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1916#[serde(rename_all = "snake_case")]
1917pub enum ProfileResolutionSource {
1918    /// `Auto` preference; ranked against the live probe.
1919    AutoRanked,
1920    /// `Pinned` preference and the profile is selectable on this probe.
1921    Pinned,
1922    /// `Pinned` preference but the profile is not selectable (unknown id or
1923    /// VRAM does not fit). Falls back to the hash profile so the call can
1924    /// still proceed deterministically.
1925    PinnedUnselectable,
1926    /// `ForceHash` preference; hash fallback selected regardless of probe.
1927    ForcedHash,
1928}
1929
1930/// Result of resolving a `ProfilePreference` against the live GPU probe.
1931#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1932pub struct ProfileResolution {
1933    pub preference: ProfilePreference,
1934    pub role: ModelRole,
1935    pub source: ProfileResolutionSource,
1936    pub profile: ModelProfile,
1937    pub selectable: bool,
1938    pub reason: String,
1939}
1940
1941/// Resolve a caller preference to a concrete profile for a given role.
1942///
1943/// Pure function — pass a synthetic `GpuProbe` for tests. The hash profile is
1944/// the guaranteed-selectable fallback for any non-`ForceHash` preference when
1945/// the pinned/auto-ranked profile is not selectable.
1946pub fn resolve_profile_preference(
1947    preference: &ProfilePreference,
1948    role: ModelRole,
1949    probe: &GpuProbe,
1950) -> ProfileResolution {
1951    let profiles = default_model_profiles();
1952    let hash_profile = profiles
1953        .iter()
1954        .find(|profile| profile.id == "tsift-local-hash-v1")
1955        .cloned()
1956        .expect("hash fallback profile is always present");
1957
1958    match preference {
1959        ProfilePreference::ForceHash => ProfileResolution {
1960            preference: preference.clone(),
1961            role,
1962            source: ProfileResolutionSource::ForcedHash,
1963            profile: hash_profile,
1964            selectable: true,
1965            reason: "caller forced the CPU/hash fallback".to_string(),
1966        },
1967        ProfilePreference::Auto => {
1968            let ranked = rank_profiles_for_role(&profiles, probe, role);
1969            let pick = ranked
1970                .iter()
1971                .find(|selection| selection.selectable)
1972                .cloned()
1973                .or_else(|| {
1974                    ranked.into_iter().next().map(|selection| ProfileSelection {
1975                        selectable: false,
1976                        ..selection
1977                    })
1978                });
1979            match pick {
1980                Some(selection) if selection.selectable => ProfileResolution {
1981                    preference: preference.clone(),
1982                    role,
1983                    source: ProfileResolutionSource::AutoRanked,
1984                    profile: selection.profile.clone(),
1985                    selectable: true,
1986                    reason: format!("auto-ranked: {}", selection.reason),
1987                },
1988                Some(selection) => ProfileResolution {
1989                    preference: preference.clone(),
1990                    role,
1991                    source: ProfileResolutionSource::AutoRanked,
1992                    profile: hash_profile,
1993                    selectable: true,
1994                    reason: format!(
1995                        "auto-ranked but no GPU profile selectable ({}); using hash fallback",
1996                        selection.reason
1997                    ),
1998                },
1999                None => ProfileResolution {
2000                    preference: preference.clone(),
2001                    role,
2002                    source: ProfileResolutionSource::AutoRanked,
2003                    profile: hash_profile,
2004                    selectable: true,
2005                    reason: "no profile matches the requested role; using hash fallback"
2006                        .to_string(),
2007                },
2008            }
2009        }
2010        ProfilePreference::Pinned(id) => match profile_by_id(id) {
2011            Some(profile) if profile.supports_role(&role) => {
2012                let selection = selection_for_profile(&profile, probe);
2013                if selection.selectable {
2014                    ProfileResolution {
2015                        preference: preference.clone(),
2016                        role,
2017                        source: ProfileResolutionSource::Pinned,
2018                        profile,
2019                        selectable: true,
2020                        reason: format!("pinned: {}", selection.reason),
2021                    }
2022                } else {
2023                    ProfileResolution {
2024                        preference: preference.clone(),
2025                        role,
2026                        source: ProfileResolutionSource::PinnedUnselectable,
2027                        profile: hash_profile,
2028                        selectable: true,
2029                        reason: format!(
2030                            "pinned {} is not selectable ({}); using hash fallback",
2031                            id, selection.reason
2032                        ),
2033                    }
2034                }
2035            }
2036            Some(_) => ProfileResolution {
2037                preference: preference.clone(),
2038                role,
2039                source: ProfileResolutionSource::PinnedUnselectable,
2040                profile: hash_profile,
2041                selectable: true,
2042                reason: format!(
2043                    "pinned {id} does not support role {:?}; using hash fallback",
2044                    role
2045                ),
2046            },
2047            None => ProfileResolution {
2048                preference: preference.clone(),
2049                role,
2050                source: ProfileResolutionSource::PinnedUnselectable,
2051                profile: hash_profile,
2052                selectable: true,
2053                reason: format!("pinned profile id {id:?} is unknown; using hash fallback"),
2054            },
2055        },
2056    }
2057}
2058
2059// ============================================================================
2060// Profile swap lifecycle (#gctrl3)
2061//
2062// `tsift local-model swap --from <id> --to <id>` is the one-command mid-run
2063// downgrade path. It combines an unload cleanup proof for the source profile
2064// with a `ProfileResolution` for the target against the post-unload probe, so
2065// a caller can decide in one step whether it is safe to load the next profile
2066// (typically qwen3-32b-q4 -> qwen3-embedding-0.6b or the hash fallback).
2067//
2068// Lease coordination stays the caller's job (they hold the holder-pid context
2069// and can chain `lease release --from` -> `swap` -> `lease acquire --to`).
2070// ============================================================================
2071
2072#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
2073#[serde(rename_all = "snake_case")]
2074pub enum SwapStatus {
2075    /// Source unload cleanup proven and target profile fits the post-unload probe.
2076    Swapped,
2077    /// Target was the CPU/hash profile; swap is always permitted once the
2078    /// source unload is proven.
2079    SwappedToHash,
2080    /// Source unload cleanup proven but the target profile does not fit the
2081    /// post-unload probe. Caller should fall back to a smaller profile or hash.
2082    UnloadProvenTargetUnselectable,
2083    /// Source unload cleanup NOT proven — caller MUST NOT load the target
2084    /// because VRAM has not been returned to baseline.
2085    UnloadNotProven,
2086    /// Source and target are the same profile id; no-op.
2087    NoOpSameProfile,
2088}
2089
2090#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
2091pub struct LocalModelSwapReport {
2092    pub from_profile_id: String,
2093    pub to_profile_id: String,
2094    pub unload: LocalModelLifecycleReport,
2095    pub target_resolution: ProfileResolution,
2096    pub swap_status: SwapStatus,
2097    pub notes: Vec<String>,
2098}
2099
2100/// Build a combined swap report: source unload lifecycle + target resolution.
2101///
2102/// Reuses `build_lifecycle_report` for the unload proof and
2103/// `resolve_profile_preference` for the target so semantics stay aligned with
2104/// the rest of the substrate.
2105#[allow(clippy::too_many_arguments)]
2106pub fn build_swap_report(
2107    from_profile: ModelProfile,
2108    to_profile: ModelProfile,
2109    pre_load_probe: GpuProbe,
2110    post_unload_probe: GpuProbe,
2111    provider_endpoint: Option<String>,
2112    provider_pid: Option<u32>,
2113    idle_ttl_seconds: u64,
2114    tolerance_mib: u64,
2115) -> LocalModelSwapReport {
2116    let unload = build_lifecycle_report(
2117        from_profile.clone(),
2118        pre_load_probe,
2119        post_unload_probe.clone(),
2120        provider_endpoint,
2121        provider_pid,
2122        idle_ttl_seconds,
2123        tolerance_mib,
2124    );
2125
2126    let target_role = to_profile
2127        .roles
2128        .first()
2129        .copied()
2130        .unwrap_or(ModelRole::Extract);
2131    let target_resolution = resolve_profile_preference(
2132        &ProfilePreference::Pinned(to_profile.id.to_string()),
2133        target_role,
2134        &post_unload_probe,
2135    );
2136
2137    let swap_status = if from_profile.id == to_profile.id {
2138        SwapStatus::NoOpSameProfile
2139    } else if !unload.cleanup.cleanup_proven {
2140        SwapStatus::UnloadNotProven
2141    } else if to_profile.concurrency == ConcurrencyClass::CpuOrHash {
2142        SwapStatus::SwappedToHash
2143    } else if target_resolution.selectable && target_resolution.profile.id == to_profile.id {
2144        SwapStatus::Swapped
2145    } else {
2146        SwapStatus::UnloadProvenTargetUnselectable
2147    };
2148
2149    let mut notes = vec![
2150        format!("swapping from {} to {}", from_profile.id, to_profile.id),
2151        format!("unload cleanup: {:?}", unload.cleanup.status),
2152        format!("target resolution: {:?}", target_resolution.source),
2153    ];
2154    if swap_status == SwapStatus::UnloadNotProven {
2155        notes.push("DO NOT load target — source unload did not prove VRAM cleanup".to_string());
2156    }
2157    if swap_status == SwapStatus::UnloadProvenTargetUnselectable {
2158        notes.push(format!(
2159            "target {} is not selectable on the post-unload probe; consider the hash fallback or a smaller profile",
2160            to_profile.id
2161        ));
2162    }
2163
2164    LocalModelSwapReport {
2165        from_profile_id: from_profile.id.to_string(),
2166        to_profile_id: to_profile.id.to_string(),
2167        unload,
2168        target_resolution,
2169        swap_status,
2170        notes,
2171    }
2172}
2173
2174#[cfg(test)]
2175mod tests {
2176    use super::*;
2177
2178    fn rtx_5090_probe() -> GpuProbe {
2179        GpuProbe {
2180            timestamp_unix_seconds: Some(1_781_000_000),
2181            available: true,
2182            gpu_name: Some("NVIDIA GeForce RTX 5090".to_string()),
2183            driver_version: Some("610.43.02".to_string()),
2184            total_vram_mib: Some(32_607),
2185            used_vram_mib: Some(179),
2186            free_vram_mib: Some(32_428),
2187            processes: Vec::new(),
2188            error: None,
2189        }
2190    }
2191
2192    fn probe_with_used_vram(used_vram_mib: u64) -> GpuProbe {
2193        let mut probe = rtx_5090_probe();
2194        probe.used_vram_mib = Some(used_vram_mib);
2195        probe.free_vram_mib = probe
2196            .total_vram_mib
2197            .map(|total| total.saturating_sub(used_vram_mib));
2198        probe
2199    }
2200
2201    #[test]
2202    fn qwen3_32b_is_default_extractor_for_clear_5090() {
2203        let report = build_status_report_with_probe(rtx_5090_probe());
2204        assert_eq!(
2205            report.recommended_extractor.as_deref(),
2206            Some("qwen3-32b-q4")
2207        );
2208        assert!(
2209            report
2210                .extractor_profiles
2211                .iter()
2212                .any(|selection| selection.profile.id == "qwen3.5-35b-a3b-q4"
2213                    && !selection.selectable)
2214        );
2215    }
2216
2217    #[test]
2218    fn hash_fallback_selects_without_gpu_probe() {
2219        let report = build_status_report_with_probe(GpuProbe::unavailable("missing"));
2220        assert_eq!(
2221            report.recommended_embedding.as_deref(),
2222            Some("tsift-local-hash-v1")
2223        );
2224        assert_eq!(report.recommended_extractor, None);
2225    }
2226
2227    #[test]
2228    fn parses_nvidia_smi_gpu_query_row() {
2229        let probe =
2230            parse_gpu_query("NVIDIA GeForce RTX 5090, 610.43.02, 32607, 179, 32428").unwrap();
2231        assert!(probe.timestamp_unix_seconds.is_some());
2232        assert_eq!(probe.gpu_name.as_deref(), Some("NVIDIA GeForce RTX 5090"));
2233        assert_eq!(probe.total_vram_mib, Some(32_607));
2234        assert_eq!(probe.used_vram_mib, Some(179));
2235        assert_eq!(probe.free_vram_mib, Some(32_428));
2236    }
2237
2238    #[test]
2239    fn lifecycle_report_plans_llamacpp_unload_and_proves_cleanup() {
2240        let profile = profile_by_id("qwen3-32b-q4").unwrap();
2241        let report = build_lifecycle_report(
2242            profile,
2243            probe_with_used_vram(200),
2244            probe_with_used_vram(820),
2245            Some("http://127.0.0.1:8080/models/unload".to_string()),
2246            Some(42),
2247            DEFAULT_IDLE_TTL_SECONDS,
2248            DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB,
2249        );
2250
2251        assert_eq!(report.lease.mode, LeaseMode::Exclusive);
2252        assert!(report.cleanup.cleanup_proven);
2253        assert_eq!(report.cleanup.status, VramCleanupStatus::Proven);
2254        assert!(report.lease.unload_actions.iter().any(|action| {
2255            action.kind == UnloadActionKind::ProviderApi
2256                && action.endpoint.as_deref() == Some("http://127.0.0.1:8080/models/unload")
2257        }));
2258        assert!(report.lease.unload_actions.iter().any(|action| {
2259            action.kind == UnloadActionKind::ProcessExit
2260                && action.command.as_ref().is_some_and(|command| {
2261                    command == &vec!["kill".to_string(), "-TERM".to_string(), "42".to_string()]
2262                })
2263        }));
2264    }
2265
2266    #[test]
2267    fn vram_cleanup_fails_when_provider_process_remains_loaded() {
2268        let pre = probe_with_used_vram(200);
2269        let mut post = probe_with_used_vram(4_000);
2270        post.processes.push(GpuProcess {
2271            pid: Some(42),
2272            process_name: "llama-server".to_string(),
2273            used_memory_mib: Some(3_000),
2274        });
2275
2276        let cleanup = evaluate_vram_cleanup(&pre, &post, DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB);
2277
2278        assert!(!cleanup.cleanup_proven);
2279        assert_eq!(cleanup.status, VramCleanupStatus::NotProven);
2280        assert_eq!(cleanup.blocking_processes.len(), 1);
2281    }
2282
2283    #[test]
2284    fn vram_cleanup_accepts_external_process_accounting() {
2285        let pre = probe_with_used_vram(200);
2286        let mut post = probe_with_used_vram(2_000);
2287        post.processes.push(GpuProcess {
2288            pid: Some(77),
2289            process_name: "python-training-job".to_string(),
2290            used_memory_mib: Some(1_600),
2291        });
2292
2293        let cleanup = evaluate_vram_cleanup(&pre, &post, DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB);
2294
2295        assert!(cleanup.cleanup_proven);
2296        assert_eq!(
2297            cleanup.status,
2298            VramCleanupStatus::ProvenByExternalAccounting
2299        );
2300        assert_eq!(cleanup.external_process_delta_mib, 1_600);
2301    }
2302
2303    #[test]
2304    fn interrupted_run_cleanup_fails_with_orphaned_provider_process() {
2305        let profile = profile_by_id("qwen3-32b-q4").unwrap();
2306        let pre = probe_with_used_vram(200);
2307        let mut post = probe_with_used_vram(8_000);
2308        post.processes.push(GpuProcess {
2309            pid: Some(1234),
2310            process_name: "ollama runner".to_string(),
2311            used_memory_mib: Some(7_000),
2312        });
2313
2314        let report = build_lifecycle_report(
2315            profile,
2316            pre,
2317            post,
2318            None,
2319            Some(1234),
2320            DEFAULT_IDLE_TTL_SECONDS,
2321            DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB,
2322        );
2323
2324        assert!(!report.cleanup.cleanup_proven);
2325        assert_eq!(report.cleanup.status, VramCleanupStatus::NotProven);
2326        assert_eq!(report.cleanup.blocking_processes.len(), 1);
2327        assert!(report.lease.unload_actions.iter().any(|action| {
2328            action.kind == UnloadActionKind::ProcessExit
2329                && action.command.as_ref().is_some_and(|command| {
2330                    command == &vec!["kill".to_string(), "-TERM".to_string(), "1234".to_string()]
2331                })
2332        }));
2333    }
2334
2335    // ---- Cooperative GPU lease registry (#gctrl1) ----
2336
2337    fn all_alive(_pid: u32) -> bool {
2338        true
2339    }
2340    fn alive_set(alive: &[u32]) -> impl Fn(u32) -> bool + '_ {
2341        move |pid| alive.contains(&pid)
2342    }
2343
2344    #[test]
2345    fn resolve_lease_file_prefers_explicit_override() {
2346        let path = resolve_lease_file(Some(Path::new("/custom/lease.json")));
2347        assert_eq!(path, PathBuf::from("/custom/lease.json"));
2348    }
2349
2350    #[test]
2351    fn resolve_lease_file_returns_env_value_when_set() {
2352        // SAFETY: env mutation is unsafe in edition 2024 because of potential
2353        // data races in multi-threaded programs. Tests run single-threaded
2354        // inside this test function and the value is restored afterwards.
2355        unsafe {
2356            std::env::set_var(LEASE_FILE_ENV_VAR, "/env/lease.json");
2357        }
2358        let path = resolve_lease_file(None);
2359        unsafe {
2360            std::env::remove_var(LEASE_FILE_ENV_VAR);
2361        }
2362        assert_eq!(path, PathBuf::from("/env/lease.json"));
2363    }
2364
2365    #[test]
2366    fn lease_registry_round_trips_through_json() {
2367        let mut registry = GpuLeaseRegistry::default();
2368        registry.leases.insert(
2369            "qwen3-32b-q4".to_string(),
2370            vec![GpuLeaseRecord {
2371                profile_id: "qwen3-32b-q4".to_string(),
2372                holder_pid: 4242,
2373                holder_command: "tsift".to_string(),
2374                acquired_at_unix_seconds: 100,
2375                lease_mode: LeaseMode::Exclusive,
2376                vram_baseline_mib: 200,
2377                idle_ttl_seconds: 0,
2378                notes: vec!["baseline".to_string()],
2379            }],
2380        );
2381        let payload = serde_json::to_string(&registry).unwrap();
2382        let back: GpuLeaseRegistry = serde_json::from_str(&payload).unwrap();
2383        assert_eq!(registry, back);
2384        assert_eq!(back.version, LEASE_REGISTRY_VERSION);
2385    }
2386
2387    #[test]
2388    fn acquire_exclusive_profile_succeeds_when_free() {
2389        let profile = profile_by_id("qwen3-32b-q4").unwrap();
2390        let mut registry = GpuLeaseRegistry::default();
2391        let acquisition = apply_acquire(
2392            &mut registry,
2393            &profile,
2394            100,
2395            "tsift",
2396            200,
2397            0,
2398            1_000,
2399            all_alive,
2400        );
2401        assert_eq!(acquisition.status, GpuLeaseAcquisitionStatus::Acquired);
2402        assert!(acquisition.conflict.is_none());
2403        assert_eq!(registry.leases["qwen3-32b-q4"].len(), 1);
2404        assert_eq!(registry.leases["qwen3-32b-q4"][0].holder_pid, 100);
2405    }
2406
2407    #[test]
2408    fn acquire_exclusive_profile_conflicts_with_live_holder() {
2409        let profile = profile_by_id("qwen3-32b-q4").unwrap();
2410        let mut registry = GpuLeaseRegistry::default();
2411        apply_acquire(
2412            &mut registry,
2413            &profile,
2414            100,
2415            "tsift",
2416            200,
2417            0,
2418            1_000,
2419            all_alive,
2420        );
2421        let second = apply_acquire(
2422            &mut registry,
2423            &profile,
2424            200,
2425            "corky",
2426            250,
2427            0,
2428            1_050,
2429            alive_set(&[100, 200]),
2430        );
2431        assert_eq!(second.status, GpuLeaseAcquisitionStatus::Conflict);
2432        let conflict = second.conflict.unwrap();
2433        assert_eq!(conflict.holder_pid, 100);
2434        assert_eq!(conflict.holder_command, "tsift");
2435        // The conflict must not overwrite the existing holder.
2436        assert_eq!(registry.leases["qwen3-32b-q4"].len(), 1);
2437        assert_eq!(registry.leases["qwen3-32b-q4"][0].holder_pid, 100);
2438    }
2439
2440    #[test]
2441    fn acquire_exclusive_profile_reclaims_when_holder_pid_dead() {
2442        let profile = profile_by_id("qwen3-32b-q4").unwrap();
2443        let mut registry = GpuLeaseRegistry::default();
2444        apply_acquire(
2445            &mut registry,
2446            &profile,
2447            100,
2448            "tsift",
2449            200,
2450            0,
2451            1_000,
2452            all_alive,
2453        );
2454        // pid 100 is gone now; only pid 200 is alive.
2455        let reclaimed = apply_acquire(
2456            &mut registry,
2457            &profile,
2458            200,
2459            "corky",
2460            250,
2461            0,
2462            1_050,
2463            alive_set(&[200]),
2464        );
2465        assert_eq!(reclaimed.status, GpuLeaseAcquisitionStatus::ReclaimedStale);
2466        assert_eq!(reclaimed.reclaimed.len(), 1);
2467        assert_eq!(registry.leases["qwen3-32b-q4"].len(), 1);
2468        assert_eq!(registry.leases["qwen3-32b-q4"][0].holder_pid, 200);
2469    }
2470
2471    #[test]
2472    fn acquire_shared_profile_allows_multiple_live_holders() {
2473        let profile = profile_by_id("qwen3-embedding-0.6b").unwrap();
2474        let mut registry = GpuLeaseRegistry::default();
2475        apply_acquire(
2476            &mut registry,
2477            &profile,
2478            100,
2479            "tsift",
2480            200,
2481            0,
2482            1_000,
2483            all_alive,
2484        );
2485        let second = apply_acquire(
2486            &mut registry,
2487            &profile,
2488            200,
2489            "headroom",
2490            250,
2491            0,
2492            1_050,
2493            alive_set(&[100, 200]),
2494        );
2495        assert_eq!(second.status, GpuLeaseAcquisitionStatus::Acquired);
2496        assert_eq!(registry.leases["qwen3-embedding-0.6b"].len(), 2);
2497    }
2498
2499    #[test]
2500    fn acquire_refreshes_when_same_holder_requests_again() {
2501        let profile = profile_by_id("qwen3-32b-q4").unwrap();
2502        let mut registry = GpuLeaseRegistry::default();
2503        apply_acquire(
2504            &mut registry,
2505            &profile,
2506            100,
2507            "tsift",
2508            200,
2509            0,
2510            1_000,
2511            all_alive,
2512        );
2513        let again = apply_acquire(
2514            &mut registry,
2515            &profile,
2516            100,
2517            "tsift",
2518            180,
2519            0,
2520            1_500,
2521            all_alive,
2522        );
2523        assert_eq!(again.status, GpuLeaseAcquisitionStatus::Refreshed);
2524        assert_eq!(registry.leases["qwen3-32b-q4"].len(), 1);
2525        assert_eq!(
2526            registry.leases["qwen3-32b-q4"][0].acquired_at_unix_seconds,
2527            1_500
2528        );
2529        assert_eq!(registry.leases["qwen3-32b-q4"][0].vram_baseline_mib, 180);
2530    }
2531
2532    #[test]
2533    fn acquire_cpu_or_hash_profile_bypasses_registry() {
2534        let profile = profile_by_id("tsift-local-hash-v1").unwrap();
2535        let mut registry = GpuLeaseRegistry::default();
2536        let bypass = apply_acquire(
2537            &mut registry,
2538            &profile,
2539            100,
2540            "tsift",
2541            0,
2542            0,
2543            1_000,
2544            all_alive,
2545        );
2546        assert_eq!(bypass.status, GpuLeaseAcquisitionStatus::CpuOrHashBypass);
2547        assert!(registry.leases.is_empty());
2548    }
2549
2550    #[test]
2551    fn idle_ttl_expires_even_when_pid_still_alive() {
2552        let profile = profile_by_id("qwen3-32b-q4").unwrap();
2553        let mut registry = GpuLeaseRegistry::default();
2554        apply_acquire(
2555            &mut registry,
2556            &profile,
2557            100,
2558            "tsift",
2559            200,
2560            60,
2561            1_000,
2562            all_alive,
2563        );
2564        // 120s later, the 60s TTL has expired; pid 100 is still alive but stale.
2565        let reclaimed = apply_acquire(
2566            &mut registry,
2567            &profile,
2568            200,
2569            "corky",
2570            250,
2571            0,
2572            1_120,
2573            all_alive,
2574        );
2575        assert_eq!(reclaimed.status, GpuLeaseAcquisitionStatus::ReclaimedStale);
2576        assert_eq!(registry.leases["qwen3-32b-q4"][0].holder_pid, 200);
2577    }
2578
2579    #[test]
2580    fn release_removes_holder_and_drops_empty_profile() {
2581        let profile = profile_by_id("qwen3-32b-q4").unwrap();
2582        let mut registry = GpuLeaseRegistry::default();
2583        apply_acquire(
2584            &mut registry,
2585            &profile,
2586            100,
2587            "tsift",
2588            200,
2589            0,
2590            1_000,
2591            all_alive,
2592        );
2593        let release = apply_release(&mut registry, "qwen3-32b-q4", 100, 1_050, all_alive);
2594        assert_eq!(release.outcome, GpuLeaseReleaseOutcome::Released);
2595        assert_eq!(release.remaining_holders, 0);
2596        assert!(registry.leases.is_empty());
2597    }
2598
2599    #[test]
2600    fn release_by_non_holder_reports_not_held() {
2601        let profile = profile_by_id("qwen3-32b-q4").unwrap();
2602        let mut registry = GpuLeaseRegistry::default();
2603        apply_acquire(
2604            &mut registry,
2605            &profile,
2606            100,
2607            "tsift",
2608            200,
2609            0,
2610            1_000,
2611            all_alive,
2612        );
2613        let release = apply_release(&mut registry, "qwen3-32b-q4", 999, 1_050, all_alive);
2614        assert_eq!(release.outcome, GpuLeaseReleaseOutcome::NotHeld);
2615        assert_eq!(registry.leases["qwen3-32b-q4"].len(), 1);
2616    }
2617
2618    #[test]
2619    fn acquire_and_release_round_trip_through_file() {
2620        let dir = tempfile_dir();
2621        let path = dir.join("gpu-lease.json");
2622        let profile = profile_by_id("qwen3-32b-q4").unwrap();
2623
2624        let mut registry = GpuLeaseRegistry::default();
2625        let acquisition = apply_acquire(
2626            &mut registry,
2627            &profile,
2628            4242,
2629            "tsift",
2630            220,
2631            0,
2632            1_000,
2633            all_alive,
2634        );
2635        assert_eq!(acquisition.status, GpuLeaseAcquisitionStatus::Acquired);
2636        write_lease_registry(&path, &registry).unwrap();
2637
2638        let read_back = read_lease_registry(&path).unwrap();
2639        assert_eq!(read_back, registry);
2640        assert_eq!(read_back.leases["qwen3-32b-q4"][0].holder_pid, 4242);
2641
2642        let release = apply_release(&mut registry, "qwen3-32b-q4", 4242, 1_050, all_alive);
2643        assert_eq!(release.outcome, GpuLeaseReleaseOutcome::Released);
2644        write_lease_registry(&path, &registry).unwrap();
2645
2646        let after = read_lease_registry(&path).unwrap();
2647        assert!(after.leases.is_empty());
2648    }
2649
2650    #[test]
2651    fn read_lease_registry_returns_default_for_missing_file() {
2652        let path = Path::new("/definitely/not/a/real/path/lease.json");
2653        let registry = read_lease_registry(path).unwrap();
2654        assert_eq!(registry, GpuLeaseRegistry::default());
2655    }
2656
2657    #[test]
2658    fn registry_lock_path_appends_lock_suffix() {
2659        assert_eq!(
2660            registry_lock_path(Path::new("/tmp/x/gpu-lease.json")),
2661            PathBuf::from("/tmp/x/gpu-lease.json.lock")
2662        );
2663    }
2664
2665    #[test]
2666    fn acquire_lease_creates_sidecar_lock_file() {
2667        let dir = tempfile_dir();
2668        let path = dir.join("gpu-lease.json");
2669        // acquire_lease runs under with_registry_lock, which opens/creates the
2670        // sidecar lock used to serialize the read-modify-write across processes.
2671        acquire_lease("qwen3-32b-q4", std::process::id(), "tsift", 0, 0, 1_000, &path).unwrap();
2672        assert!(
2673            registry_lock_path(&path).exists(),
2674            "sidecar lock file should exist after a locked acquire"
2675        );
2676    }
2677
2678    #[test]
2679    fn apply_renew_slides_heartbeat_so_ttl_holder_survives() {
2680        let mut registry = GpuLeaseRegistry::default();
2681        let profile = profile_by_id("qwen3-32b-q4").unwrap();
2682        // Acquire with a 100s idle TTL at t=1_000.
2683        apply_acquire(&mut registry, &profile, 100, "tsift", 200, 100, 1_000, all_alive);
2684        // Heartbeat at t=1_050 (still within the TTL window) slides the anchor.
2685        let renew = apply_renew(&mut registry, "qwen3-32b-q4", 100, 1_050, all_alive);
2686        assert_eq!(renew.outcome, GpuLeaseRenewOutcome::Renewed);
2687        assert_eq!(renew.renewed_at_unix_seconds, Some(1_050));
2688        assert_eq!(
2689            registry.leases["qwen3-32b-q4"][0].acquired_at_unix_seconds,
2690            1_050
2691        );
2692        // At t=1_120 the age since the heartbeat (1_050) is 70s < 100s TTL, so
2693        // it survives — whereas without the renewal (anchor 1_000) it would have
2694        // expired at 1_100.
2695        let pruned = prune_stale_leases(&mut registry, 1_120, all_alive);
2696        assert!(pruned.is_empty());
2697        assert!(registry.leases.contains_key("qwen3-32b-q4"));
2698    }
2699
2700    #[test]
2701    fn apply_renew_reports_profile_absent_for_unheld_profile() {
2702        let mut registry = GpuLeaseRegistry::default();
2703        let renew = apply_renew(&mut registry, "qwen3-32b-q4", 100, 1_000, all_alive);
2704        assert_eq!(renew.outcome, GpuLeaseRenewOutcome::ProfileAbsent);
2705        assert!(renew.renewed_at_unix_seconds.is_none());
2706    }
2707
2708    #[test]
2709    fn reap_leases_reclaims_dead_pid_and_reports_emptied_profile() {
2710        let dir = tempfile_dir();
2711        let path = dir.join("gpu-lease.json");
2712        let mut registry = GpuLeaseRegistry::default();
2713        registry.leases.insert(
2714            "qwen3-32b-q4".to_string(),
2715            vec![GpuLeaseRecord {
2716                profile_id: "qwen3-32b-q4".to_string(),
2717                // A pid far above any live process — `kill -0` reports it dead,
2718                // simulating a session that crashed without releasing.
2719                holder_pid: 4_000_000_000,
2720                holder_command: "crashed-session".to_string(),
2721                acquired_at_unix_seconds: 1_000,
2722                lease_mode: LeaseMode::Exclusive,
2723                vram_baseline_mib: 200,
2724                idle_ttl_seconds: 0,
2725                notes: Vec::new(),
2726            }],
2727        );
2728        write_lease_registry(&path, &registry).unwrap();
2729
2730        let reap = reap_leases(2_000, &path).unwrap();
2731        assert_eq!(reap.reclaimed.len(), 1);
2732        assert_eq!(reap.emptied_profiles, vec!["qwen3-32b-q4".to_string()]);
2733        let after = read_lease_registry(&path).unwrap();
2734        assert!(after.leases.is_empty());
2735    }
2736
2737    #[test]
2738    fn renew_lease_round_trips_through_file() {
2739        let dir = tempfile_dir();
2740        let path = dir.join("gpu-lease.json");
2741        let pid = std::process::id();
2742        // ttl=0 → no TTL staleness; the live pid keeps the lease, so the renew
2743        // exercises the file round-trip + timestamp slide without TTL timing.
2744        acquire_lease("qwen3-32b-q4", pid, "tsift", 0, 0, 1_000, &path).unwrap();
2745        let renew = renew_lease("qwen3-32b-q4", pid, 5_000, &path).unwrap();
2746        assert_eq!(renew.outcome, GpuLeaseRenewOutcome::Renewed);
2747        let registry = read_lease_registry(&path).unwrap();
2748        assert_eq!(
2749            registry.leases["qwen3-32b-q4"][0].acquired_at_unix_seconds,
2750            5_000
2751        );
2752    }
2753
2754    #[test]
2755    fn prune_stale_leaves_healthy_entries_alone() {
2756        let mut registry = GpuLeaseRegistry::default();
2757        registry.leases.insert(
2758            "qwen3-32b-q4".to_string(),
2759            vec![GpuLeaseRecord {
2760                profile_id: "qwen3-32b-q4".to_string(),
2761                holder_pid: 100,
2762                holder_command: "tsift".to_string(),
2763                acquired_at_unix_seconds: 1_000,
2764                lease_mode: LeaseMode::Exclusive,
2765                vram_baseline_mib: 200,
2766                idle_ttl_seconds: 0,
2767                notes: Vec::new(),
2768            }],
2769        );
2770        let pruned = prune_stale_leases(&mut registry, 1_010, alive_set(&[100]));
2771        assert!(pruned.is_empty());
2772        assert!(registry.leases.contains_key("qwen3-32b-q4"));
2773    }
2774
2775    /// Serialize tests that mutate the shared process endpoint env vars so they
2776    /// do not race under parallel `cargo test`.
2777    fn env_lock() -> std::sync::MutexGuard<'static, ()> {
2778        static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
2779        ENV_LOCK
2780            .lock()
2781            .unwrap_or_else(|poisoned| poisoned.into_inner())
2782    }
2783
2784    fn tempfile_dir() -> PathBuf {
2785        use std::sync::atomic::{AtomicU64, Ordering};
2786        static COUNTER: AtomicU64 = AtomicU64::new(0);
2787        // Per-test uniqueness: `current_unix_seconds()` collides when tests run
2788        // within the same second in parallel, so include a monotonic counter.
2789        let dir = std::env::temp_dir().join(format!(
2790            "tsift-lease-test-{}-{}-{}",
2791            std::process::id(),
2792            current_unix_seconds(),
2793            COUNTER.fetch_add(1, Ordering::Relaxed)
2794        ));
2795        std::fs::create_dir_all(&dir).unwrap();
2796        dir
2797    }
2798
2799    // ---- Per-call profile preference (#gctrl2) ----
2800
2801    #[test]
2802    fn profile_preference_parses_cli_value() {
2803        assert_eq!(ProfilePreference::from_cli(None), ProfilePreference::Auto);
2804        assert_eq!(
2805            ProfilePreference::from_cli(Some("")),
2806            ProfilePreference::Auto
2807        );
2808        assert_eq!(
2809            ProfilePreference::from_cli(Some("hash")),
2810            ProfilePreference::ForceHash
2811        );
2812        assert_eq!(
2813            ProfilePreference::from_cli(Some("tsift-local-hash-v1")),
2814            ProfilePreference::ForceHash
2815        );
2816        assert_eq!(
2817            ProfilePreference::from_cli(Some("qwen3-32b-q4")),
2818            ProfilePreference::Pinned("qwen3-32b-q4".to_string())
2819        );
2820    }
2821
2822    #[test]
2823    fn resolve_auto_picks_recommended_gpu_profile_on_clear_5090() {
2824        let probe = rtx_5090_probe();
2825        let resolution =
2826            resolve_profile_preference(&ProfilePreference::Auto, ModelRole::Extract, &probe);
2827        assert_eq!(resolution.source, ProfileResolutionSource::AutoRanked);
2828        assert!(resolution.selectable);
2829        assert_eq!(resolution.profile.id, "qwen3-32b-q4");
2830    }
2831
2832    #[test]
2833    fn resolve_auto_falls_back_to_hash_when_gpu_unavailable() {
2834        let probe = GpuProbe::unavailable("missing");
2835        let resolution =
2836            resolve_profile_preference(&ProfilePreference::Auto, ModelRole::Extract, &probe);
2837        assert_eq!(resolution.source, ProfileResolutionSource::AutoRanked);
2838        assert_eq!(resolution.profile.id, "tsift-local-hash-v1");
2839        assert!(resolution.selectable);
2840        assert!(resolution.reason.contains("no GPU profile selectable"));
2841    }
2842
2843    #[test]
2844    fn resolve_pinned_selectable_profile_is_used_as_is() {
2845        let probe = rtx_5090_probe();
2846        let resolution = resolve_profile_preference(
2847            &ProfilePreference::Pinned("qwen3-embedding-0.6b".to_string()),
2848            ModelRole::Embed,
2849            &probe,
2850        );
2851        assert_eq!(resolution.source, ProfileResolutionSource::Pinned);
2852        assert_eq!(resolution.profile.id, "qwen3-embedding-0.6b");
2853        assert!(resolution.selectable);
2854    }
2855
2856    #[test]
2857    fn resolve_pinned_profile_with_wrong_role_falls_back_to_hash() {
2858        let probe = rtx_5090_probe();
2859        let resolution = resolve_profile_preference(
2860            &ProfilePreference::Pinned("qwen3-embedding-0.6b".to_string()),
2861            ModelRole::Extract,
2862            &probe,
2863        );
2864        assert_eq!(
2865            resolution.source,
2866            ProfileResolutionSource::PinnedUnselectable
2867        );
2868        assert_eq!(resolution.profile.id, "tsift-local-hash-v1");
2869        assert!(resolution.reason.contains("does not support role"));
2870    }
2871
2872    #[test]
2873    fn resolve_pinned_profile_that_does_not_fit_vram_falls_back_to_hash() {
2874        // 30 GiB used → only ~2.6 GiB free → qwen3-32b-q4 (~28 GiB) won't fit.
2875        let probe = probe_with_used_vram(30_000);
2876        let resolution = resolve_profile_preference(
2877            &ProfilePreference::Pinned("qwen3-32b-q4".to_string()),
2878            ModelRole::Extract,
2879            &probe,
2880        );
2881        assert_eq!(
2882            resolution.source,
2883            ProfileResolutionSource::PinnedUnselectable
2884        );
2885        assert_eq!(resolution.profile.id, "tsift-local-hash-v1");
2886        assert!(resolution.selectable);
2887        assert!(resolution.reason.contains("not selectable"));
2888    }
2889
2890    #[test]
2891    fn resolve_force_hash_always_uses_hash_profile() {
2892        let probe = rtx_5090_probe();
2893        let resolution =
2894            resolve_profile_preference(&ProfilePreference::ForceHash, ModelRole::Extract, &probe);
2895        assert_eq!(resolution.source, ProfileResolutionSource::ForcedHash);
2896        assert_eq!(resolution.profile.id, "tsift-local-hash-v1");
2897        assert!(resolution.selectable);
2898        assert!(resolution.reason.contains("forced"));
2899    }
2900
2901    #[test]
2902    fn resolve_pinned_unknown_profile_id_falls_back_to_hash() {
2903        let probe = rtx_5090_probe();
2904        let resolution = resolve_profile_preference(
2905            &ProfilePreference::Pinned("not-a-real-profile".to_string()),
2906            ModelRole::Embed,
2907            &probe,
2908        );
2909        assert_eq!(
2910            resolution.source,
2911            ProfileResolutionSource::PinnedUnselectable
2912        );
2913        assert_eq!(resolution.profile.id, "tsift-local-hash-v1");
2914        assert!(resolution.reason.contains("unknown"));
2915    }
2916
2917    // ---- Provider endpoint configurability (#portconf) ----
2918
2919    #[test]
2920    fn resolve_endpoint_returns_explicit_override_for_any_strategy() {
2921        for strategy in [
2922            UnloadStrategy::LlamaCppRouterUnload,
2923            UnloadStrategy::OllamaKeepAliveZero,
2924            UnloadStrategy::VllmSleep,
2925            UnloadStrategy::ProcessExit,
2926            UnloadStrategy::None,
2927        ] {
2928            let resolved = resolve_provider_endpoint(&strategy, Some("http://custom:9999/path"));
2929            assert_eq!(
2930                resolved, "http://custom:9999/path",
2931                "explicit override should win for {strategy:?}"
2932            );
2933        }
2934    }
2935
2936    #[test]
2937    fn resolve_endpoint_uses_compile_time_default_when_no_env_no_explicit() {
2938        // Serialize against sibling env-mutating tests (parallel `cargo test`).
2939        let _env = env_lock();
2940        // SAFETY: env-mutating tests are serialized via `env_lock`; the vars are
2941        // cleared before returning.
2942        unsafe {
2943            std::env::remove_var(LLAMA_CPP_ENDPOINT_ENV_VAR);
2944            std::env::remove_var(OLLAMA_ENDPOINT_ENV_VAR);
2945            std::env::remove_var(VLLM_ENDPOINT_ENV_VAR);
2946        }
2947        assert_eq!(
2948            resolve_provider_endpoint(&UnloadStrategy::LlamaCppRouterUnload, None),
2949            DEFAULT_LLAMA_CPP_ENDPOINT
2950        );
2951        assert_eq!(
2952            resolve_provider_endpoint(&UnloadStrategy::OllamaKeepAliveZero, None),
2953            DEFAULT_OLLAMA_ENDPOINT
2954        );
2955        assert_eq!(
2956            resolve_provider_endpoint(&UnloadStrategy::VllmSleep, None),
2957            DEFAULT_VLLM_ENDPOINT
2958        );
2959        assert_eq!(
2960            resolve_provider_endpoint(&UnloadStrategy::ProcessExit, None),
2961            ""
2962        );
2963        assert_eq!(resolve_provider_endpoint(&UnloadStrategy::None, None), "");
2964    }
2965
2966    #[test]
2967    fn resolve_endpoint_env_var_overrides_default_for_llama_cpp() {
2968        let _env = env_lock();
2969        // SAFETY: see note in the previous test.
2970        unsafe {
2971            std::env::set_var(
2972                LLAMA_CPP_ENDPOINT_ENV_VAR,
2973                "http://127.0.0.1:8081/models/unload",
2974            );
2975        }
2976        let resolved = resolve_provider_endpoint(&UnloadStrategy::LlamaCppRouterUnload, None);
2977        unsafe {
2978            std::env::remove_var(LLAMA_CPP_ENDPOINT_ENV_VAR);
2979        }
2980        assert_eq!(resolved, "http://127.0.0.1:8081/models/unload");
2981    }
2982
2983    #[test]
2984    fn resolve_endpoint_blank_env_var_falls_back_to_default() {
2985        let _env = env_lock();
2986        // SAFETY: see note above.
2987        unsafe {
2988            std::env::set_var(LLAMA_CPP_ENDPOINT_ENV_VAR, "   ");
2989        }
2990        let resolved = resolve_provider_endpoint(&UnloadStrategy::LlamaCppRouterUnload, None);
2991        unsafe {
2992            std::env::remove_var(LLAMA_CPP_ENDPOINT_ENV_VAR);
2993        }
2994        assert_eq!(resolved, DEFAULT_LLAMA_CPP_ENDPOINT);
2995    }
2996
2997    #[test]
2998    fn build_unload_actions_picks_up_env_var_for_llama_cpp_endpoint() {
2999        let _env = env_lock();
3000        let profile = profile_by_id("qwen3-32b-q4").unwrap();
3001        // SAFETY: see note above.
3002        unsafe {
3003            std::env::set_var(
3004                LLAMA_CPP_ENDPOINT_ENV_VAR,
3005                "http://127.0.0.1:8081/models/unload",
3006            );
3007        }
3008        let actions = build_unload_actions(&profile, None, Some(42));
3009        unsafe {
3010            std::env::remove_var(LLAMA_CPP_ENDPOINT_ENV_VAR);
3011        }
3012        let unload_action = actions
3013            .iter()
3014            .find(|action| action.kind == UnloadActionKind::ProviderApi)
3015            .expect("provider api action present");
3016        assert_eq!(
3017            unload_action.endpoint.as_deref(),
3018            Some("http://127.0.0.1:8081/models/unload")
3019        );
3020    }
3021
3022    // ---- Profile swap lifecycle (#gctrl3) ----
3023
3024    fn probe_pair(pre_used: u64, post_used: u64) -> (GpuProbe, GpuProbe) {
3025        (
3026            probe_with_used_vram(pre_used),
3027            probe_with_used_vram(post_used),
3028        )
3029    }
3030
3031    #[test]
3032    fn swap_to_same_profile_is_noop() {
3033        let from = profile_by_id("qwen3-32b-q4").unwrap();
3034        let to = profile_by_id("qwen3-32b-q4").unwrap();
3035        let (pre, post) = probe_pair(200, 200);
3036        let report = build_swap_report(
3037            from,
3038            to,
3039            pre,
3040            post,
3041            None,
3042            None,
3043            DEFAULT_IDLE_TTL_SECONDS,
3044            DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB,
3045        );
3046        assert_eq!(report.swap_status, SwapStatus::NoOpSameProfile);
3047    }
3048
3049    #[test]
3050    fn swap_from_big_to_small_embedding_when_cleanup_proven_is_swapped() {
3051        let from = profile_by_id("qwen3-32b-q4").unwrap();
3052        let to = profile_by_id("qwen3-embedding-0.6b").unwrap();
3053        // Source was using ~28 GiB; after unload it returns to ~200 MiB.
3054        let (pre, post) = probe_pair(28_000, 200);
3055        let report = build_swap_report(
3056            from,
3057            to,
3058            pre,
3059            post,
3060            None,
3061            Some(42),
3062            DEFAULT_IDLE_TTL_SECONDS,
3063            DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB,
3064        );
3065        assert_eq!(report.swap_status, SwapStatus::Swapped);
3066        assert!(report.unload.cleanup.cleanup_proven);
3067        assert_eq!(report.target_resolution.profile.id, "qwen3-embedding-0.6b");
3068    }
3069
3070    #[test]
3071    fn swap_to_hash_fallback_is_swapped_to_hash_when_cleanup_proven() {
3072        let from = profile_by_id("qwen3-32b-q4").unwrap();
3073        let to = profile_by_id("tsift-local-hash-v1").unwrap();
3074        let (pre, post) = probe_pair(28_000, 200);
3075        let report = build_swap_report(
3076            from,
3077            to,
3078            pre,
3079            post,
3080            None,
3081            Some(42),
3082            DEFAULT_IDLE_TTL_SECONDS,
3083            DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB,
3084        );
3085        assert_eq!(report.swap_status, SwapStatus::SwappedToHash);
3086        assert!(report.unload.cleanup.cleanup_proven);
3087    }
3088
3089    #[test]
3090    fn swap_blocks_when_source_unload_not_proven() {
3091        let from = profile_by_id("qwen3-32b-q4").unwrap();
3092        let to = profile_by_id("qwen3-embedding-0.6b").unwrap();
3093        // Baseline VRAM is ~200 MiB before load. Orphaned llama-server process
3094        // holds ~7 GiB after "unload", so cleanup is NOT proven.
3095        let pre = probe_with_used_vram(200);
3096        let mut post = probe_with_used_vram(8_000);
3097        post.processes.push(GpuProcess {
3098            pid: Some(42),
3099            process_name: "llama-server".to_string(),
3100            used_memory_mib: Some(7_000),
3101        });
3102        let report = build_swap_report(
3103            from,
3104            to,
3105            pre,
3106            post,
3107            None,
3108            Some(42),
3109            DEFAULT_IDLE_TTL_SECONDS,
3110            DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB,
3111        );
3112        assert_eq!(report.swap_status, SwapStatus::UnloadNotProven);
3113        assert!(!report.unload.cleanup.cleanup_proven);
3114        assert!(
3115            report
3116                .notes
3117                .iter()
3118                .any(|note| note.contains("DO NOT load target"))
3119        );
3120    }
3121
3122    #[test]
3123    fn swap_reports_target_unselectable_when_post_unload_vram_still_high() {
3124        let from = profile_by_id("qwen3-32b-q4").unwrap();
3125        let to = profile_by_id("qwen3-32b-q4").unwrap();
3126        // Source is qwen3-32b-q4 itself; after unload, only ~3 GiB free — the
3127        // target 32B footprint (28.7 GiB) does not fit. Cleanup is proven
3128        // (post <= pre + tolerance), but the target cannot reload.
3129        let pre = probe_with_used_vram(29_500);
3130        let post = probe_with_used_vram(29_600);
3131        let report = build_swap_report(
3132            from,
3133            to,
3134            pre,
3135            post,
3136            None,
3137            None,
3138            DEFAULT_IDLE_TTL_SECONDS,
3139            DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB,
3140        );
3141        // from == to is the NoOpSameProfile path; pick distinct ids instead.
3142        assert_eq!(report.swap_status, SwapStatus::NoOpSameProfile);
3143        // Re-run with a distinct target to exercise UnloadProvenTargetUnselectable.
3144        let from = profile_by_id("qwen3-32b-q4").unwrap();
3145        let to = profile_by_id("qwen3-embedding-8b").unwrap();
3146        let pre = probe_with_used_vram(30_000);
3147        let post = probe_with_used_vram(30_500);
3148        let report = build_swap_report(
3149            from,
3150            to,
3151            pre,
3152            post,
3153            None,
3154            None,
3155            DEFAULT_IDLE_TTL_SECONDS,
3156            DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB,
3157        );
3158        assert_eq!(
3159            report.swap_status,
3160            SwapStatus::UnloadProvenTargetUnselectable
3161        );
3162        assert!(
3163            report
3164                .notes
3165                .iter()
3166                .any(|note| note.contains("not selectable on the post-unload probe"))
3167        );
3168    }
3169
3170    // =========================================================================
3171    // #kgunloadpost: build_unload_actions owns execution — pure helper tests
3172    // =========================================================================
3173
3174    #[test]
3175    fn rewrite_unload_body_model_replaces_model_field() {
3176        let original = r#"{"model":"qwen3-32b-q4-ollama-default","prompt":"","keep_alive":0}"#;
3177        let rewritten = rewrite_unload_body_model(original, "hf.co/Qwen/Qwen3-32B-GGUF:Q4_K_M");
3178        let value: serde_json::Value =
3179            serde_json::from_str(&rewritten).expect("rewritten body is valid JSON");
3180        assert_eq!(
3181            value["model"].as_str(),
3182            Some("hf.co/Qwen/Qwen3-32B-GGUF:Q4_K_M")
3183        );
3184        // Other fields preserved.
3185        assert_eq!(value["keep_alive"].as_i64(), Some(0));
3186        assert_eq!(value["prompt"].as_str(), Some(""));
3187    }
3188
3189    #[test]
3190    fn rewrite_unload_body_model_preserves_body_when_override_is_empty() {
3191        let original = r#"{"model":"default-tag","keep_alive":0}"#;
3192        let rewritten = rewrite_unload_body_model(original, "");
3193        // Empty override must not blank out the model field — falls through.
3194        assert_eq!(rewritten, original);
3195    }
3196
3197    #[test]
3198    fn rewrite_unload_body_model_falls_back_on_invalid_json() {
3199        let original = "not valid json {{{";
3200        let rewritten = rewrite_unload_body_model(original, "any-tag");
3201        assert_eq!(rewritten, original);
3202    }
3203
3204    #[test]
3205    fn normalize_unload_url_appends_generate_path_for_bare_host() {
3206        let url = normalize_unload_url("http://127.0.0.1:11434");
3207        assert_eq!(url, "http://127.0.0.1:11434/api/generate");
3208    }
3209
3210    #[test]
3211    fn normalize_unload_url_idempotent_for_full_generate_url() {
3212        let url = normalize_unload_url("http://127.0.0.1:11434/api/generate");
3213        assert_eq!(url, "http://127.0.0.1:11434/api/generate");
3214    }
3215
3216    #[test]
3217    fn normalize_unload_url_strips_trailing_slash() {
3218        let url = normalize_unload_url("http://127.0.0.1:11434/");
3219        assert_eq!(url, "http://127.0.0.1:11434/api/generate");
3220    }
3221
3222    #[test]
3223    fn prepare_unload_request_returns_none_for_non_api_actions() {
3224        let noop = ProviderUnloadAction {
3225            kind: UnloadActionKind::Noop,
3226            label: "noop".to_string(),
3227            command: None,
3228            http_method: None,
3229            endpoint: None,
3230            body_json: None,
3231            required: false,
3232        };
3233        assert!(prepare_unload_request(&noop, "any-tag").is_none());
3234    }
3235
3236    #[test]
3237    fn prepare_unload_request_applies_model_override_to_body() {
3238        // Mirrors the OllamaKeepAliveZero action shape produced by
3239        // build_unload_actions: body carries the profile's model_ref, and the
3240        // resolved override must replace it (the bug fixed by #kgunloadpost).
3241        let action = ProviderUnloadAction {
3242            kind: UnloadActionKind::ProviderApi,
3243            label: "ollama keep_alive zero".to_string(),
3244            command: Some(vec!["ollama".to_string(), "stop".to_string()]),
3245            http_method: Some("POST".to_string()),
3246            endpoint: Some("http://127.0.0.1:11434".to_string()),
3247            body_json: Some(
3248                r#"{"model":"profile-default-tag","prompt":"","keep_alive":0}"#.to_string(),
3249            ),
3250            required: true,
3251        };
3252        let req = prepare_unload_request(&action, "override-tag")
3253            .expect("ProviderApi action prepares a request");
3254        assert_eq!(req.url, "http://127.0.0.1:11434/api/generate");
3255        assert!(req.body.contains("\"model\":\"override-tag\""));
3256        assert!(!req.body.contains("profile-default-tag"));
3257        assert_eq!(
3258            req.fallback_command,
3259            Some(vec!["ollama".to_string(), "stop".to_string()])
3260        );
3261    }
3262
3263    #[test]
3264    fn prepare_unload_request_synthesizes_body_when_plan_has_none() {
3265        let action = ProviderUnloadAction {
3266            kind: UnloadActionKind::ProviderApi,
3267            label: "synthesized".to_string(),
3268            command: None,
3269            http_method: Some("POST".to_string()),
3270            endpoint: Some("http://host:11434".to_string()),
3271            body_json: None,
3272            required: true,
3273        };
3274        let req = prepare_unload_request(&action, "synth-tag").unwrap();
3275        assert!(req.body.contains("\"model\":\"synth-tag\""));
3276        assert!(req.body.contains("\"keep_alive\":0"));
3277    }
3278
3279    #[test]
3280    fn execute_unload_actions_reports_non_api_as_skipped() {
3281        let actions = vec![ProviderUnloadAction {
3282            kind: UnloadActionKind::Noop,
3283            label: "no GPU unload required".to_string(),
3284            command: None,
3285            http_method: None,
3286            endpoint: None,
3287            body_json: None,
3288            required: false,
3289        }];
3290        let results = execute_unload_actions(&actions, "any-tag");
3291        assert_eq!(results.len(), 1);
3292        assert!(!results[0].executed);
3293        assert_eq!(results[0].outcome, "skipped: non-API action");
3294    }
3295}