Blame for src/status.rs

//! 本机 GPU 状态诊断：读 Ollama `/api/ps` 的真实 VRAM 分配。
//!
//! # 为什么要读 `/api/ps` 而不是猜
//!
//! 大多数应用会通过 `num_gpu` 参数"控制"GPU 层数——**这是错的**。
//! Ollama 自己会按模型结构和显存选择 Metal/CUDA 后端与层数；
//! 强行传 `num_gpu` 会触发 runner 重载，反而更慢。
//!
//! 正确的做法：**让 Ollama 自己决定，然后如实报告它实际分配了多少显存。**

use serde::{Deserialize, Serialize};

/// 本机 Ollama 的默认地址。
pub const DEFAULT_OLLAMA_BASE_URL: &str = "http://127.0.0.1:11434";

/// 单个模型的 GPU 驻留状态。
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct GpuModelStatus {
    /// 模型名（如 `qwen3-06b:latest`）。
    pub name: String,
    /// 模型总大小（字节）。
    pub model_bytes: u64,
    /// 实际分配在 VRAM 的字节数（`size_vram`）。
    pub gpu_memory_bytes: u64,
    /// 上下文长度（Ollama 新版才返回）。
    pub context_length: Option<u64>,
    /// 驻留过期时间（keep_alive 到期）。
    pub expires_at: Option<String>,
    /// GPU 驻留百分比（0~100，近似值）。
    /// 这是**驻留估算**，不是 GPU 利用率，也不是 tokens/s。
    pub gpu_residency_percent: u8,
}

/// 本机 GPU 加速总状态。
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct GpuAccelerationStatus {
    /// Ollama 是否在线。
    pub online: bool,
    /// 查询时间（毫秒时间戳）。
    pub checked_at_ms: i64,
    /// 已加载的模型列表。
    pub models: Vec<GpuModelStatus>,
}

#[derive(Debug, Default, Deserialize)]
struct PsResponse {
    #[serde(default)]
    models: Vec<PsModel>,
}

#[derive(Debug, Default, Deserialize)]
struct PsModel {
    #[serde(default)]
    name: String,
    #[serde(default)]
    size: u64,
    #[serde(default, rename = "size_vram")]
    size_vram: u64,
    #[serde(default, rename = "context_length")]
    context_length: Option<u64>,
    #[serde(default, rename = "expires_at")]
    expires_at: Option<String>,
}

/// 计算 GPU 驻留百分比（纯函数，可单测）。
pub fn residency_percent(size: u64, gpu: u64) -> u8 {
    if size == 0 {
        return 0;
    }
    ((gpu.min(size) as u128 * 100 / size as u128).min(100)) as u8
}

/// 把 `/api/ps` 的原始响应转成结构化状态（纯函数，可单测）。
pub fn from_ps_json(value: &str, now_ms: i64) -> Result<GpuAccelerationStatus, String> {
    let parsed: PsResponse =
        serde_json::from_str(value).map_err(|e| format!("无法解析 /api/ps 响应：{e}"))?;
    let models = parsed
        .models
        .into_iter()
        .map(|model| GpuModelStatus {
            gpu_residency_percent: residency_percent(model.size, model.size_vram),
            name: model.name,
            model_bytes: model.size,
            gpu_memory_bytes: model.size_vram,
            context_length: model.context_length,
            expires_at: model.expires_at,
        })
        .collect();
    Ok(GpuAccelerationStatus {
        online: true,
        checked_at_ms: now_ms,
        models,
    })
}

/// 查询本机 Ollama 的 GPU 状态（HTTP 请求）。
pub async fn query_status(base_url: Option<&str>) -> Result<GpuAccelerationStatus, String> {
    let base = base_url
        .unwrap_or(DEFAULT_OLLAMA_BASE_URL)
        .trim_end_matches('/');
    let client = reqwest::Client::builder()
        .connect_timeout(std::time::Duration::from_secs(1))
        .build()
        .map_err(|e| format!("初始化状态查询失败：{e}"))?;
    let response = client
        .get(format!("{base}/api/ps"))
        .timeout(std::time::Duration::from_secs(3))
        .send()
        .await
        .map_err(|e| format!("无法读取本机 Ollama GPU 状态：{e}"))?;
    if !response.status().is_success() {
        return Err(format!("本机 Ollama 状态接口返回 {}", response.status()));
    }
    let text = response
        .text()
        .await
        .map_err(|e| format!("无法读取响应体：{e}"))?;
    from_ps_json(&text, now_ms())
}

fn now_ms() -> i64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_millis() as i64)
        .unwrap_or(0)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn reports_real_residency_and_handles_cpu_fallback() {
        let status = from_ps_json(
            r#"{"models":[
                {"name":"gpu-model","size":1000,"size_vram":1000,"context_length":4096,"expires_at":"soon"},
                {"name":"partial-model","size":1000,"size_vram":500},
                {"name":"cpu-model","size":1000,"size_vram":0},
                {"name":"unknown-size","size":0,"size_vram":0}
            ]}"#,
            0,
        )
        .unwrap();
        assert!(status.online);
        assert_eq!(status.models[0].gpu_residency_percent, 100);
        assert_eq!(status.models[1].gpu_residency_percent, 50);
        assert_eq!(status.models[2].gpu_residency_percent, 0);
        assert_eq!(status.models[3].gpu_residency_percent, 0);
        assert_eq!(status.models[0].context_length, Some(4096));
        // 序列化字段名（camelCase）
        let wire = serde_json::to_value(&status).unwrap();
        assert_eq!(wire["models"][0]["gpuResidencyPercent"], 100);
        assert_eq!(wire["models"][0]["gpuMemoryBytes"], 1000);
    }

    #[test]
    fn accepts_no_loaded_models() {
        let status = from_ps_json("{}", 0).unwrap();
        assert!(status.online);
        assert!(status.models.is_empty());
    }
}