From ac0155d48e48c64ca400435a67fddf303d186493 Mon Sep 17 00:00:00 2001 From: CrimsonMartin Date: Fri, 28 Aug 2026 21:09:18 +0000 Subject: [PATCH 1/6] feat: port MCP server (stdio) onto upstream base MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Re-applies the fork's MCP (Model Context Protocol) server — JSON-RPC 2.0 over stdio with generate/chat/embed tools and a fox://models resource — on top of the upstream tree. Adapted to the current engine API: submit_request's SubmitError result is surfaced as a JSON-RPC error, sampling uses the centralized defaults with SamplingParams::default(), and registry setup goes through the new RegistryConfig::embedded() constructor (mirrors ServeArgs defaults) shared with tests. --- src/cli/mcp.rs | 54 +++ src/cli/mod.rs | 5 + src/lib.rs | 1 + src/mcp/mod.rs | 806 +++++++++++++++++++++++++++++++++++ src/model_registry/config.rs | 43 ++ 5 files changed, 909 insertions(+) create mode 100644 src/cli/mcp.rs create mode 100644 src/mcp/mod.rs diff --git a/src/cli/mcp.rs b/src/cli/mcp.rs new file mode 100644 index 0000000..8bcca9d --- /dev/null +++ b/src/cli/mcp.rs @@ -0,0 +1,54 @@ +// `fox mcp` — start an MCP (Model Context Protocol) server over stdio. +// Designed for IDE integration (Cursor, VS Code, Claude Code). + +use std::path::PathBuf; +use std::sync::Arc; + +use anyhow::Result; +use clap::Parser; + +use crate::mcp::McpServer; +use crate::model_registry::{ModelRegistry, RegistryConfig}; + +use super::get_gpu_memory_bytes; +use super::load_aliases; +use super::models_dir as default_models_dir; + +#[derive(Parser, Debug)] +pub struct McpArgs { + /// Path to GGUF model file (optional; models are loaded on demand) + #[arg(long, env = "FOX_MODEL_PATH")] + pub model_path: Option, + + /// Path to aliases TOML file. Default: ~/.config/ferrumox/aliases.toml + #[arg(long, env = "FOX_ALIAS_FILE")] + pub alias_file: Option, +} + +pub async fn run_mcp(args: McpArgs) -> Result<()> { + // Log to stderr only — stdout is the MCP transport. + tracing_subscriber::fmt() + .with_writer(std::io::stderr) + .with_env_filter( + tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("ferrumox=info,warn")), + ) + .init(); + + let gpu_memory_bytes = get_gpu_memory_bytes(); + let aliases = load_aliases(args.alias_file); + let models_dir = default_models_dir(); + + let registry_cfg = RegistryConfig::embedded(models_dir.clone(), gpu_memory_bytes); + let registry = Arc::new(ModelRegistry::new(registry_cfg, aliases)); + + if let Some(path) = &args.model_path { + eprintln!("mcp: pre-loading model from {:?}", path); + registry + .get_or_load(path.to_string_lossy().as_ref()) + .await?; + } + + eprintln!("fox mcp server listening on stdio"); + McpServer::new(registry, models_dir).run().await +} diff --git a/src/cli/mod.rs b/src/cli/mod.rs index 4441d3a..61f68c6 100644 --- a/src/cli/mod.rs +++ b/src/cli/mod.rs @@ -6,6 +6,7 @@ pub mod bench_kv; pub mod bench_prefill; pub mod bench_spec; pub mod list; +pub mod mcp; pub mod models; pub mod probe; pub mod ps; @@ -67,6 +68,8 @@ pub enum Command { Search(search::SearchArgs), /// Manage model name aliases Alias(alias::AliasArgs), + /// Start an MCP (Model Context Protocol) server over stdio for IDE integration + Mcp(mcp::McpArgs), } /// Known subcommand names — anything else is treated as `fox run `. @@ -86,6 +89,7 @@ const SUBCOMMANDS: &[&str] = &[ "models", "search", "alias", + "mcp", "help", ]; @@ -123,5 +127,6 @@ pub async fn run() -> anyhow::Result<()> { Command::Models(args) => models::run_models(args).await, Command::Search(args) => search::run_search(args).await, Command::Alias(args) => alias::run_alias(args).await, + Command::Mcp(args) => mcp::run_mcp(args).await, } } diff --git a/src/lib.rs b/src/lib.rs index b3346f6..01041f5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -5,6 +5,7 @@ pub mod cli; pub mod config; pub(crate) mod engine; pub(crate) mod kv_cache; +pub(crate) mod mcp; pub mod metrics; pub mod model_registry; pub mod registry; diff --git a/src/mcp/mod.rs b/src/mcp/mod.rs new file mode 100644 index 0000000..8994113 --- /dev/null +++ b/src/mcp/mod.rs @@ -0,0 +1,806 @@ +// MCP (Model Context Protocol) server — JSON-RPC 2.0 over stdio transport. +// Reads Content-Length framed messages from stdin, writes responses to stdout. +// All diagnostic logging goes to stderr (stdout is reserved for the protocol). + +use std::io::{BufRead, Write}; +use std::sync::Arc; + +use serde::{Deserialize, Serialize}; + +use crate::model_registry::ModelRegistry; + +// --------------------------------------------------------------------------- +// JSON-RPC 2.0 types +// --------------------------------------------------------------------------- + +#[derive(Debug, Deserialize)] +pub(crate) struct JsonRpcRequest { + #[allow(dead_code)] + pub jsonrpc: String, + pub id: Option, + pub method: String, + #[serde(default)] + pub params: serde_json::Value, +} + +#[derive(Debug, Serialize)] +pub(crate) struct JsonRpcResponse { + pub jsonrpc: &'static str, + #[serde(skip_serializing_if = "Option::is_none")] + pub id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub result: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub error: Option, +} + +#[derive(Debug, Serialize)] +pub(crate) struct JsonRpcError { + pub code: i64, + pub message: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub data: Option, +} + +impl JsonRpcResponse { + fn success(id: Option, result: serde_json::Value) -> Self { + Self { + jsonrpc: "2.0", + id, + result: Some(result), + error: None, + } + } + + fn error(id: Option, code: i64, message: impl Into) -> Self { + Self { + jsonrpc: "2.0", + id, + result: None, + error: Some(JsonRpcError { + code, + message: message.into(), + data: None, + }), + } + } +} + +// JSON-RPC error codes +const METHOD_NOT_FOUND: i64 = -32601; +const INVALID_PARAMS: i64 = -32602; +const INTERNAL_ERROR: i64 = -32603; + +// --------------------------------------------------------------------------- +// Content-Length framing (LSP-style) +// --------------------------------------------------------------------------- + +/// Read one Content-Length framed JSON message from the reader. +/// Returns `None` on EOF. +pub(crate) fn read_message(reader: &mut impl BufRead) -> std::io::Result> { + let mut content_length: Option = None; + let mut header_line = String::new(); + + loop { + header_line.clear(); + let n = reader.read_line(&mut header_line)?; + if n == 0 { + return Ok(None); // EOF + } + let trimmed = header_line.trim(); + if trimmed.is_empty() { + break; // end of headers + } + if let Some(val) = trimmed.strip_prefix("Content-Length:") { + if let Ok(len) = val.trim().parse::() { + content_length = Some(len); + } + } + } + + let len = match content_length { + Some(l) => l, + None => return Ok(None), + }; + + let mut body = vec![0u8; len]; + reader.read_exact(&mut body)?; + Ok(Some(String::from_utf8_lossy(&body).into_owned())) +} + +/// Write one Content-Length framed JSON message to the writer. +pub(crate) fn write_message(writer: &mut impl Write, body: &str) -> std::io::Result<()> { + write!(writer, "Content-Length: {}\r\n\r\n{}", body.len(), body)?; + writer.flush() +} + +// --------------------------------------------------------------------------- +// MCP tool/resource definitions +// --------------------------------------------------------------------------- + +fn tool_definitions() -> serde_json::Value { + serde_json::json!([ + { + "name": "generate", + "description": "Generate text completion from a prompt.", + "inputSchema": { + "type": "object", + "properties": { + "model": { "type": "string", "description": "Model name or path" }, + "prompt": { "type": "string", "description": "Text prompt" }, + "max_tokens": { "type": "integer", "description": "Maximum tokens to generate" }, + "temperature": { "type": "number", "description": "Sampling temperature" } + }, + "required": ["model", "prompt"] + } + }, + { + "name": "chat", + "description": "Chat completion with message history.", + "inputSchema": { + "type": "object", + "properties": { + "model": { "type": "string", "description": "Model name or path" }, + "messages": { + "type": "array", + "items": { + "type": "object", + "properties": { + "role": { "type": "string" }, + "content": { "type": "string" } + }, + "required": ["role", "content"] + }, + "description": "Chat messages" + }, + "max_tokens": { "type": "integer", "description": "Maximum tokens to generate" }, + "temperature": { "type": "number", "description": "Sampling temperature" } + }, + "required": ["model", "messages"] + } + }, + { + "name": "embed", + "description": "Generate embeddings for text input.", + "inputSchema": { + "type": "object", + "properties": { + "model": { "type": "string", "description": "Model name or path" }, + "input": { + "description": "Text or array of texts to embed", + "oneOf": [ + { "type": "string" }, + { "type": "array", "items": { "type": "string" } } + ] + } + }, + "required": ["model", "input"] + } + } + ]) +} + +fn resource_definitions() -> serde_json::Value { + serde_json::json!([ + { + "uri": "fox://models", + "name": "models", + "description": "List available models.", + "mimeType": "application/json" + } + ]) +} + +// --------------------------------------------------------------------------- +// MCP server +// --------------------------------------------------------------------------- + +pub(crate) struct McpServer { + registry: Arc, + models_dir: std::path::PathBuf, +} + +impl McpServer { + pub fn new(registry: Arc, models_dir: std::path::PathBuf) -> Self { + Self { + registry, + models_dir, + } + } + + /// Run the stdio event loop. Blocks until stdin is closed. + pub async fn run(self) -> anyhow::Result<()> { + let server = Arc::new(self); + + // Spawn a blocking reader thread since stdin is synchronous. + let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::(); + std::thread::spawn(move || { + let stdin = std::io::stdin(); + let mut reader = stdin.lock(); + loop { + match read_message(&mut reader) { + Ok(Some(msg)) => { + if tx.send(msg).is_err() { + break; + } + } + Ok(None) => break, // EOF + Err(e) => { + eprintln!("mcp: stdin read error: {e}"); + break; + } + } + } + }); + + while let Some(msg) = rx.recv().await { + let server = server.clone(); + let response = server.handle_message(&msg).await; + if let Some(resp) = response { + let body = serde_json::to_string(&resp).unwrap_or_default(); + let stdout = std::io::stdout(); + let mut out = stdout.lock(); + if let Err(e) = write_message(&mut out, &body) { + eprintln!("mcp: stdout write error: {e}"); + break; + } + } + } + + Ok(()) + } + + async fn handle_message(&self, raw: &str) -> Option { + let req: JsonRpcRequest = match serde_json::from_str(raw) { + Ok(r) => r, + Err(e) => { + return Some(JsonRpcResponse::error( + None, + -32700, // Parse error + format!("invalid JSON: {e}"), + )); + } + }; + + // Notifications (no id) do not get responses. + req.id.as_ref()?; + + let result = match req.method.as_str() { + "initialize" => self.handle_initialize(), + "initialized" => return None, // notification acknowledgment + "tools/list" => self.handle_tools_list(), + "tools/call" => self.handle_tools_call(&req.params).await, + "resources/list" => self.handle_resources_list(), + "resources/read" => self.handle_resources_read(&req.params), + _ => Err((METHOD_NOT_FOUND, format!("unknown method: {}", req.method))), + }; + + Some(match result { + Ok(val) => JsonRpcResponse::success(req.id, val), + Err((code, msg)) => JsonRpcResponse::error(req.id, code, msg), + }) + } + + fn handle_initialize(&self) -> Result { + Ok(serde_json::json!({ + "protocolVersion": "2024-11-05", + "capabilities": { + "tools": {}, + "resources": {} + }, + "serverInfo": { + "name": "fox", + "version": env!("CARGO_PKG_VERSION") + } + })) + } + + fn handle_tools_list(&self) -> Result { + Ok(serde_json::json!({ "tools": tool_definitions() })) + } + + fn handle_resources_list(&self) -> Result { + Ok(serde_json::json!({ "resources": resource_definitions() })) + } + + fn handle_resources_read( + &self, + params: &serde_json::Value, + ) -> Result { + let uri = params + .get("uri") + .and_then(|v| v.as_str()) + .ok_or_else(|| (INVALID_PARAMS, "missing uri".to_string()))?; + + match uri { + "fox://models" => { + let models = crate::cli::list_models(&self.models_dir).unwrap_or_default(); + let list: Vec = models + .iter() + .map(|(path, meta)| { + let name = path + .file_stem() + .and_then(|s| s.to_str()) + .unwrap_or("unknown"); + serde_json::json!({ + "name": name, + "size": meta.len(), + }) + }) + .collect(); + + Ok(serde_json::json!({ + "contents": [{ + "uri": "fox://models", + "mimeType": "application/json", + "text": serde_json::to_string(&list).unwrap_or_else(|_| "[]".to_string()) + }] + })) + } + _ => Err((INVALID_PARAMS, format!("unknown resource uri: {uri}"))), + } + } + + async fn handle_tools_call( + &self, + params: &serde_json::Value, + ) -> Result { + let name = params + .get("name") + .and_then(|v| v.as_str()) + .ok_or_else(|| (INVALID_PARAMS, "missing tool name".to_string()))?; + let arguments = params + .get("arguments") + .cloned() + .unwrap_or(serde_json::Value::Object(serde_json::Map::new())); + + match name { + "generate" => self.tool_generate(&arguments).await, + "chat" => self.tool_chat(&arguments).await, + "embed" => self.tool_embed(&arguments).await, + _ => Err((INVALID_PARAMS, format!("unknown tool: {name}"))), + } + } + + // ----------------------------------------------------------------------- + // Tool implementations + // ----------------------------------------------------------------------- + + async fn tool_generate( + &self, + args: &serde_json::Value, + ) -> Result { + let model = args + .get("model") + .and_then(|v| v.as_str()) + .ok_or_else(|| (INVALID_PARAMS, "missing model".to_string()))?; + let prompt = args + .get("prompt") + .and_then(|v| v.as_str()) + .ok_or_else(|| (INVALID_PARAMS, "missing prompt".to_string()))?; + let max_tokens = args + .get("max_tokens") + .and_then(|v| v.as_u64()) + .unwrap_or(512) as usize; + let temperature = args + .get("temperature") + .and_then(|v| v.as_f64()) + .unwrap_or(crate::api::shared::sampling_defaults::TEMPERATURE as f64) + as f32; + + let entry = self + .registry + .get_or_load(model) + .await + .map_err(|e| (INTERNAL_ERROR, format!("failed to load model: {e}")))?; + let engine = &entry.engine; + + let messages = vec![("user".to_string(), prompt.to_string())]; + let prompt_text = engine.apply_chat_template(&messages).unwrap_or_else(|_| { + messages + .iter() + .map(|(r, c)| format!("{r}: {c}")) + .collect::>() + .join("\n") + }); + + let tokens = engine + .tokenize(&prompt_text) + .map_err(|e| (INTERNAL_ERROR, format!("tokenize failed: {e}")))?; + + let text = run_inference(engine, tokens, max_tokens, temperature).await?; + + Ok(mcp_text_result(&text)) + } + + async fn tool_chat( + &self, + args: &serde_json::Value, + ) -> Result { + let model = args + .get("model") + .and_then(|v| v.as_str()) + .ok_or_else(|| (INVALID_PARAMS, "missing model".to_string()))?; + let messages_val = args + .get("messages") + .and_then(|v| v.as_array()) + .ok_or_else(|| (INVALID_PARAMS, "missing messages array".to_string()))?; + let max_tokens = args + .get("max_tokens") + .and_then(|v| v.as_u64()) + .unwrap_or(512) as usize; + let temperature = args + .get("temperature") + .and_then(|v| v.as_f64()) + .unwrap_or(crate::api::shared::sampling_defaults::TEMPERATURE as f64) + as f32; + + let messages: Vec<(String, String)> = messages_val + .iter() + .filter_map(|m| { + let role = m.get("role")?.as_str()?.to_string(); + let content = m.get("content")?.as_str()?.to_string(); + Some((role, content)) + }) + .collect(); + + if messages.is_empty() { + return Err((INVALID_PARAMS, "messages array is empty".to_string())); + } + + let entry = self + .registry + .get_or_load(model) + .await + .map_err(|e| (INTERNAL_ERROR, format!("failed to load model: {e}")))?; + let engine = &entry.engine; + + let prompt_text = engine.apply_chat_template(&messages).unwrap_or_else(|_| { + messages + .iter() + .map(|(r, c)| format!("{r}: {c}")) + .collect::>() + .join("\n") + }); + + let tokens = engine + .tokenize(&prompt_text) + .map_err(|e| (INTERNAL_ERROR, format!("tokenize failed: {e}")))?; + + let text = run_inference(engine, tokens, max_tokens, temperature).await?; + + Ok(mcp_text_result(&text)) + } + + async fn tool_embed( + &self, + args: &serde_json::Value, + ) -> Result { + let model = args + .get("model") + .and_then(|v| v.as_str()) + .ok_or_else(|| (INVALID_PARAMS, "missing model".to_string()))?; + let input = args + .get("input") + .ok_or_else(|| (INVALID_PARAMS, "missing input".to_string()))?; + + let texts: Vec = if let Some(s) = input.as_str() { + vec![s.to_string()] + } else if let Some(arr) = input.as_array() { + arr.iter() + .filter_map(|v| v.as_str().map(String::from)) + .collect() + } else { + return Err(( + INVALID_PARAMS, + "input must be a string or array".to_string(), + )); + }; + + let entry = self + .registry + .get_or_load(model) + .await + .map_err(|e| (INTERNAL_ERROR, format!("failed to load model: {e}")))?; + let engine = &entry.engine; + + let mut embeddings = Vec::with_capacity(texts.len()); + for text in &texts { + let emb = engine + .embed(text) + .await + .map_err(|e| (INTERNAL_ERROR, format!("embedding failed: {e}")))?; + embeddings.push(emb); + } + + let result = serde_json::to_string(&embeddings).unwrap_or_else(|_| "[]".to_string()); + + Ok(mcp_text_result(&result)) + } +} + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +fn mcp_text_result(text: &str) -> serde_json::Value { + serde_json::json!({ + "content": [{ + "type": "text", + "text": text + }] + }) +} + +async fn run_inference( + engine: &crate::engine::InferenceEngine, + tokens: Vec, + max_tokens: usize, + temperature: f32, +) -> Result { + use crate::api::shared::sampling_defaults; + use crate::scheduler::{InferenceRequest, SamplingParams}; + + let sampling = SamplingParams { + temperature, + top_p: sampling_defaults::TOP_P, + ..Default::default() + }; + + let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); + let req_id = engine.next_request_id(); + let req = InferenceRequest::new(req_id, tokens, max_tokens, sampling, tx); + engine + .submit_request(req) + .map_err(|e| (INTERNAL_ERROR, format!("request rejected: {e}")))?; + + let mut result = String::new(); + while let Some(token) = rx.recv().await { + if !token.text.is_empty() { + result.push_str(&token.text); + } + if token.stop_reason.is_some() { + break; + } + } + + Ok(result) +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + + /// A server over an empty registry in a temp models dir — enough for every + /// protocol-level test; nothing here loads a real model. + fn test_server(dir: &std::path::Path) -> McpServer { + let mut cfg = + crate::model_registry::RegistryConfig::embedded(dir.to_path_buf(), 4 * 1024 * 1024); + cfg.max_context_len = Some(512); + cfg.keep_alive_secs = 0; + let registry = Arc::new(ModelRegistry::new(cfg, std::collections::HashMap::new())); + McpServer::new(registry, dir.to_path_buf()) + } + + #[test] + fn test_read_message_basic() { + let input = b"Content-Length: 17\r\n\r\n{\"hello\":\"world\"}"; + let mut cursor = std::io::Cursor::new(input.as_slice()); + let msg = read_message(&mut cursor).unwrap(); + assert_eq!(msg, Some("{\"hello\":\"world\"}".to_string())); + } + + #[test] + fn test_read_message_eof() { + let input = b""; + let mut cursor = std::io::Cursor::new(input.as_slice()); + let msg = read_message(&mut cursor).unwrap(); + assert!(msg.is_none()); + } + + #[test] + fn test_write_message() { + let mut buf = Vec::new(); + write_message(&mut buf, "{\"ok\":true}").unwrap(); + let output = String::from_utf8(buf).unwrap(); + assert!(output.starts_with("Content-Length: 11\r\n\r\n")); + assert!(output.ends_with("{\"ok\":true}")); + } + + #[test] + fn test_roundtrip_message() { + let body = r#"{"jsonrpc":"2.0","id":1,"method":"initialize"}"#; + let mut buf = Vec::new(); + write_message(&mut buf, body).unwrap(); + + let mut cursor = std::io::Cursor::new(buf.as_slice()); + let read_back = read_message(&mut cursor).unwrap(); + assert_eq!(read_back, Some(body.to_string())); + } + + #[test] + fn test_json_rpc_response_success_serialization() { + let resp = JsonRpcResponse::success(Some(serde_json::json!(1)), serde_json::json!("ok")); + let json = serde_json::to_value(&resp).unwrap(); + assert_eq!(json["jsonrpc"], "2.0"); + assert_eq!(json["id"], 1); + assert_eq!(json["result"], "ok"); + assert!(json.get("error").is_none()); + } + + #[test] + fn test_json_rpc_response_error_serialization() { + let resp = JsonRpcResponse::error(Some(serde_json::json!(2)), -32600, "invalid request"); + let json = serde_json::to_value(&resp).unwrap(); + assert_eq!(json["jsonrpc"], "2.0"); + assert_eq!(json["id"], 2); + assert_eq!(json["error"]["code"], -32600); + assert_eq!(json["error"]["message"], "invalid request"); + assert!(json.get("result").is_none()); + } + + #[test] + fn test_json_rpc_request_deserialization() { + let raw = r#"{"jsonrpc":"2.0","id":1,"method":"tools/list"}"#; + let req: JsonRpcRequest = serde_json::from_str(raw).unwrap(); + assert_eq!(req.method, "tools/list"); + assert_eq!(req.id, Some(serde_json::json!(1))); + } + + #[test] + fn test_json_rpc_request_with_params() { + let raw = r#"{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"generate","arguments":{"model":"llama","prompt":"hi"}}}"#; + let req: JsonRpcRequest = serde_json::from_str(raw).unwrap(); + assert_eq!(req.method, "tools/call"); + assert_eq!(req.params["name"], "generate"); + assert_eq!(req.params["arguments"]["prompt"], "hi"); + } + + #[test] + fn test_tool_definitions_has_all_tools() { + let tools = tool_definitions(); + let arr = tools.as_array().unwrap(); + let names: Vec<&str> = arr.iter().map(|t| t["name"].as_str().unwrap()).collect(); + assert!(names.contains(&"generate")); + assert!(names.contains(&"chat")); + assert!(names.contains(&"embed")); + assert_eq!(names.len(), 3); + } + + #[test] + fn test_resource_definitions_has_models() { + let resources = resource_definitions(); + let arr = resources.as_array().unwrap(); + assert_eq!(arr.len(), 1); + assert_eq!(arr[0]["uri"], "fox://models"); + } + + #[test] + fn test_mcp_text_result_format() { + let result = mcp_text_result("hello world"); + let content = result["content"].as_array().unwrap(); + assert_eq!(content.len(), 1); + assert_eq!(content[0]["type"], "text"); + assert_eq!(content[0]["text"], "hello world"); + } + + #[tokio::test] + async fn test_handle_initialize() { + let dir = tempfile::tempdir().unwrap(); + let server = test_server(dir.path()); + + let raw = r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}"#; + let resp = server.handle_message(raw).await.unwrap(); + assert!(resp.error.is_none()); + let result = resp.result.unwrap(); + assert_eq!(result["serverInfo"]["name"], "fox"); + assert!(result["capabilities"]["tools"].is_object()); + assert!(result["capabilities"]["resources"].is_object()); + } + + #[tokio::test] + async fn test_handle_tools_list() { + let dir = tempfile::tempdir().unwrap(); + let server = test_server(dir.path()); + + let raw = r#"{"jsonrpc":"2.0","id":2,"method":"tools/list"}"#; + let resp = server.handle_message(raw).await.unwrap(); + assert!(resp.error.is_none()); + let tools = resp.result.unwrap()["tools"].as_array().unwrap().clone(); + assert_eq!(tools.len(), 3); + } + + #[tokio::test] + async fn test_handle_resources_list() { + let dir = tempfile::tempdir().unwrap(); + let server = test_server(dir.path()); + + let raw = r#"{"jsonrpc":"2.0","id":3,"method":"resources/list"}"#; + let resp = server.handle_message(raw).await.unwrap(); + assert!(resp.error.is_none()); + let resources = resp.result.unwrap()["resources"] + .as_array() + .unwrap() + .clone(); + assert_eq!(resources.len(), 1); + assert_eq!(resources[0]["uri"], "fox://models"); + } + + #[tokio::test] + async fn test_handle_resources_read_models() { + let dir = tempfile::tempdir().unwrap(); + // Create a fake model file + std::fs::write(dir.path().join("test-model.gguf"), b"fake").unwrap(); + + let server = test_server(dir.path()); + + let raw = + r#"{"jsonrpc":"2.0","id":4,"method":"resources/read","params":{"uri":"fox://models"}}"#; + let resp = server.handle_message(raw).await.unwrap(); + assert!(resp.error.is_none()); + let result = resp.result.unwrap(); + let contents = result["contents"].as_array().unwrap(); + assert_eq!(contents[0]["uri"], "fox://models"); + let text = contents[0]["text"].as_str().unwrap(); + let models: Vec = serde_json::from_str(text).unwrap(); + assert_eq!(models.len(), 1); + assert_eq!(models[0]["name"], "test-model"); + } + + #[tokio::test] + async fn test_handle_unknown_method() { + let dir = tempfile::tempdir().unwrap(); + let server = test_server(dir.path()); + + let raw = r#"{"jsonrpc":"2.0","id":5,"method":"nonexistent/method"}"#; + let resp = server.handle_message(raw).await.unwrap(); + assert!(resp.error.is_some()); + assert_eq!(resp.error.unwrap().code, METHOD_NOT_FOUND); + } + + #[tokio::test] + async fn test_handle_invalid_json() { + let dir = tempfile::tempdir().unwrap(); + let server = test_server(dir.path()); + + let resp = server.handle_message("not valid json").await.unwrap(); + assert!(resp.error.is_some()); + assert_eq!(resp.error.unwrap().code, -32700); + } + + #[tokio::test] + async fn test_handle_notification_no_response() { + let dir = tempfile::tempdir().unwrap(); + let server = test_server(dir.path()); + + let raw = r#"{"jsonrpc":"2.0","method":"initialized"}"#; + let resp = server.handle_message(raw).await; + assert!(resp.is_none()); + } + + #[tokio::test] + async fn test_handle_resources_read_unknown_uri() { + let dir = tempfile::tempdir().unwrap(); + let server = test_server(dir.path()); + + let raw = r#"{"jsonrpc":"2.0","id":6,"method":"resources/read","params":{"uri":"fox://unknown"}}"#; + let resp = server.handle_message(raw).await.unwrap(); + assert!(resp.error.is_some()); + assert_eq!(resp.error.unwrap().code, INVALID_PARAMS); + } + + #[tokio::test] + async fn test_handle_tools_call_unknown_tool() { + let dir = tempfile::tempdir().unwrap(); + let server = test_server(dir.path()); + + let raw = r#"{"jsonrpc":"2.0","id":7,"method":"tools/call","params":{"name":"nonexistent","arguments":{}}}"#; + let resp = server.handle_message(raw).await.unwrap(); + assert!(resp.error.is_some()); + assert_eq!(resp.error.unwrap().code, INVALID_PARAMS); + } +} diff --git a/src/model_registry/config.rs b/src/model_registry/config.rs index f2543b7..ac19737 100644 --- a/src/model_registry/config.rs +++ b/src/model_registry/config.rs @@ -111,3 +111,46 @@ pub struct RegistryConfig { /// Useful for MoE models (e.g. DeepSeek, Mixtral) where expert weights don't fit in VRAM. pub moe_offload_cpu: bool, } + +impl RegistryConfig { + /// Config for embedded entry points that host a registry without going through + /// `fox serve` (e.g. `fox mcp`, tests). Values mirror `ServeArgs`' flag defaults; + /// when a default changes there, change it here too. + pub fn embedded(models_dir: PathBuf, gpu_memory_bytes: usize) -> Self { + Self { + models_dir, + max_models: 1, + max_batch_size: 32, + max_queue_depth: 0, + max_prefill_chunk: 512, + rs_rollback: 4, + context_shift: true, + context_keep: 0, + reranking: false, + cache_ram_bytes: 0, + kv_reuse: true, + slot_prompt_similarity: 0.1, + speculative: false, + spec_ngram: 2, + spec_draft_len: 4, + draft_model: None, + mmproj: None, + mtp_model: None, + lora_modules: Vec::new(), + primary_model: None, + max_context_len: None, + block_size: 16, + gpu_memory_bytes, + gpu_memory_fraction: 0.85, + metrics: None, + keep_alive_secs: 300, + type_k: kv_type::F16, + type_v: kv_type::F16, + n_gpu_layers: -1, + main_gpu: 0, + split_mode: 1, + tensor_split: Vec::new(), + moe_offload_cpu: false, + } + } +} From 8a7f1208ac16fc32bc2e6c6c44971e8fbb2e8c8b Mon Sep 17 00:00:00 2001 From: CrimsonMartin Date: Fri, 28 Aug 2026 21:11:22 +0000 Subject: [PATCH 2/6] feat: port multi-path model discovery and gpu-info CLI Re-applies the fork's `fox discover` (scans ferrumox, HuggingFace, Ollama and LM Studio model directories plus FOX_MODEL_DIRS for GGUF files, with shard grouping and dedup) and `fox gpu-info` (CUDA/Metal backend, VRAM and driver diagnostics) on the upstream base. Both were self-contained; only CLI wiring changed. Adds the walkdir dependency. --- Cargo.lock | 29 ++ Cargo.toml | 1 + src/cli/discover.rs | 82 +++++ src/cli/gpu_info.rs | 304 ++++++++++++++++ src/cli/mod.rs | 10 + src/lib.rs | 1 + src/model_discovery.rs | 791 +++++++++++++++++++++++++++++++++++++++++ 7 files changed, 1218 insertions(+) create mode 100644 src/cli/discover.rs create mode 100644 src/cli/gpu_info.rs create mode 100644 src/model_discovery.rs diff --git a/Cargo.lock b/Cargo.lock index 8d00330..d01abdb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -569,6 +569,7 @@ dependencies = [ "tracing", "tracing-subscriber", "uuid", + "walkdir", ] [[package]] @@ -1782,6 +1783,15 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "same-file" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +dependencies = [ + "winapi-util", +] + [[package]] name = "scopeguard" version = "1.2.0" @@ -2444,6 +2454,16 @@ version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" +[[package]] +name = "walkdir" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" +dependencies = [ + "same-file", + "winapi-util", +] + [[package]] name = "want" version = "0.3.1" @@ -2640,6 +2660,15 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" +[[package]] +name = "winapi-util" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" +dependencies = [ + "windows-sys 0.61.2", +] + [[package]] name = "winapi-x86_64-pc-windows-gnu" version = "0.4.0" diff --git a/Cargo.toml b/Cargo.toml index 8fee794..e136188 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -51,6 +51,7 @@ dashmap = "6" toml = "0.8" bytes = "1" dirs = "5" +walkdir = "2" # `json` is NOT a default feature, and without it the `tojson` filter is unknown. # Every native tool-use chat template (Qwen, Hermes, Mistral) renders its tool # listing with `{{ tool | tojson }}`, so leaving it out makes those templates fail diff --git a/src/cli/discover.rs b/src/cli/discover.rs new file mode 100644 index 0000000..3b8f51f --- /dev/null +++ b/src/cli/discover.rs @@ -0,0 +1,82 @@ +// `fox discover` — scan well-known directories for GGUF models. + +use std::path::{Path, PathBuf}; + +use anyhow::Result; +use clap::Parser; +use crossterm::style::Color; + +use super::theme; +use crate::cli::format_size; +use crate::model_discovery::{discover_models, parse_model_dirs}; + +#[derive(Parser, Debug)] +pub struct DiscoverArgs { + /// Semicolon-separated list of additional directories to scan for GGUF models + #[arg(long, env = "FOX_MODEL_DIRS")] + pub model_dirs: Option, +} + +pub async fn run_discover(args: DiscoverArgs) -> Result<()> { + let extra: Vec = args + .model_dirs + .as_deref() + .map(parse_model_dirs) + .unwrap_or_default(); + + let models = discover_models(&extra); + + if models.is_empty() { + eprintln!("No models found. Run `fox pull ` to download one."); + return Ok(()); + } + + println!("Found {} models:\n", models.len()); + + let max_name = models + .iter() + .map(|m| m.name.len()) + .max() + .unwrap_or(4) + .max(4); + let name_col = max_name + 2; + + theme::print_table_header(&[ + ("NAME", name_col), + ("SIZE", 10), + ("SOURCE", 14), + ("PATH", 20), + ]); + theme::print_separator(name_col + 46); + + for m in &models { + print!("{: String { + if let Some(home) = dirs::home_dir() { + if let Ok(rest) = path.strip_prefix(&home) { + return format!("~/{}", rest.display()); + } + } + path.display().to_string() +} diff --git a/src/cli/gpu_info.rs b/src/cli/gpu_info.rs new file mode 100644 index 0000000..744565e --- /dev/null +++ b/src/cli/gpu_info.rs @@ -0,0 +1,304 @@ +// `fox gpu-info` — display GPU backend, VRAM, and driver details. + +use anyhow::Result; +use clap::Parser; + +#[derive(Parser, Debug)] +pub struct GpuInfoArgs {} + +pub async fn run_gpu_info(_args: GpuInfoArgs) -> Result<()> { + println!("GPU Information"); + + #[cfg(fox_stub)] + { + println!(" Backend: CPU only (stub build)"); + println!(" No GPU acceleration detected"); + } + + #[cfg(not(fox_stub))] + { + if let Some(info) = detect_cuda() { + print_cuda_info(&info); + } else if cfg!(target_os = "macos") { + print_metal_info(); + } else { + println!(" Backend: CPU only"); + println!(" No GPU acceleration detected"); + } + } + + println!(); + println!("Build Info"); + print_build_info(); + + Ok(()) +} + +#[cfg(not(fox_stub))] +struct CudaInfo { + devices: Vec, + driver_version: Option, + cuda_version: Option, +} + +#[cfg(not(fox_stub))] +struct CudaDevice { + name: String, + memory_total_mib: usize, + memory_free_mib: usize, + compute_cap: Option, +} + +#[cfg(not(fox_stub))] +fn detect_cuda() -> Option { + let nvidia_smi = if cfg!(target_os = "windows") { + "nvidia-smi.exe" + } else { + "nvidia-smi" + }; + + let out = std::process::Command::new(nvidia_smi) + .args([ + "--query-gpu=name,memory.total,memory.free,driver_version,compute_cap", + "--format=csv,noheader,nounits", + ]) + .output() + .ok()?; + + if !out.status.success() { + return None; + } + + let stdout = std::str::from_utf8(&out.stdout).ok()?.trim().to_string(); + if stdout.is_empty() { + return None; + } + + let mut devices = Vec::new(); + let mut driver_version = None; + + for line in stdout.lines() { + let parts: Vec<&str> = line.splitn(5, ',').collect(); + if parts.len() < 4 { + continue; + } + let name = parts[0].trim().to_string(); + let total_mib: usize = parts[1].trim().parse().unwrap_or(0); + let free_mib: usize = parts[2].trim().parse().unwrap_or(0); + let drv = parts[3].trim().to_string(); + let compute = parts.get(4).map(|s| s.trim().to_string()); + + if driver_version.is_none() && !drv.is_empty() { + driver_version = Some(drv); + } + + devices.push(CudaDevice { + name, + memory_total_mib: total_mib, + memory_free_mib: free_mib, + compute_cap: compute.filter(|s| !s.is_empty()), + }); + } + + if devices.is_empty() { + return None; + } + + let cuda_version = detect_cuda_version(nvidia_smi); + + Some(CudaInfo { + devices, + driver_version, + cuda_version, + }) +} + +#[cfg(not(fox_stub))] +fn detect_cuda_version(nvidia_smi: &str) -> Option { + let out = std::process::Command::new(nvidia_smi).output().ok()?; + if !out.status.success() { + return None; + } + let stdout = std::str::from_utf8(&out.stdout).ok()?; + for line in stdout.lines() { + if let Some(idx) = line.find("CUDA Version:") { + let rest = line[idx + "CUDA Version:".len()..].trim(); + let ver = rest.split_whitespace().next().unwrap_or(rest); + return Some(ver.to_string()); + } + } + None +} + +#[cfg(not(fox_stub))] +fn print_cuda_info(info: &CudaInfo) { + println!(" Backend: CUDA"); + for (i, dev) in info.devices.iter().enumerate() { + println!(" Device {}: {}", i, dev.name); + println!(" VRAM Total: {} MB", dev.memory_total_mib); + println!(" VRAM Free: {} MB", dev.memory_free_mib); + if let Some(ref cap) = dev.compute_cap { + println!(" Compute Cap: {}", cap); + } + } + if let Some(ref drv) = info.driver_version { + println!(" Driver Version: {}", drv); + } + if let Some(ref ver) = info.cuda_version { + println!(" CUDA Version: {}", ver); + } +} + +#[cfg(all(not(fox_stub), target_os = "macos"))] +fn print_metal_info() { + println!(" Backend: Metal"); + + if let Some(gpu_name) = detect_macos_gpu_name() { + println!(" Device 0: {}", gpu_name); + } + if let Some(mem_mb) = detect_macos_memory_mb() { + println!(" Unified Memory: {} MB", mem_mb); + } + if let Some(metal_ver) = detect_macos_metal_version() { + println!(" Metal Version: {}", metal_ver); + } +} + +#[cfg(all(not(fox_stub), not(target_os = "macos")))] +fn print_metal_info() { + println!(" Backend: CPU only"); + println!(" No GPU acceleration detected"); +} + +#[cfg(all(not(fox_stub), target_os = "macos"))] +fn detect_macos_gpu_name() -> Option { + let out = std::process::Command::new("system_profiler") + .args(["SPDisplaysDataType", "-json"]) + .output() + .ok()?; + if !out.status.success() { + return None; + } + let json: serde_json::Value = serde_json::from_slice(&out.stdout).ok()?; + let displays = json.get("SPDisplaysDataType")?.as_array()?; + let first = displays.first()?; + first + .get("sppci_model") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()) +} + +#[cfg(all(not(fox_stub), target_os = "macos"))] +fn detect_macos_memory_mb() -> Option { + let out = std::process::Command::new("sysctl") + .args(["-n", "hw.memsize"]) + .output() + .ok()?; + if !out.status.success() { + return None; + } + let bytes: usize = std::str::from_utf8(&out.stdout).ok()?.trim().parse().ok()?; + Some(bytes / (1024 * 1024)) +} + +#[cfg(all(not(fox_stub), target_os = "macos"))] +fn detect_macos_metal_version() -> Option { + let out = std::process::Command::new("system_profiler") + .args(["SPDisplaysDataType", "-json"]) + .output() + .ok()?; + if !out.status.success() { + return None; + } + let json: serde_json::Value = serde_json::from_slice(&out.stdout).ok()?; + let displays = json.get("SPDisplaysDataType")?.as_array()?; + let first = displays.first()?; + first + .get("spdisplays_metal_supported") + .or_else(|| first.get("spdisplays_mtlgpufamilysupport")) + .and_then(|v| v.as_str()) + .map(|s| s.to_string()) +} + +fn print_build_info() { + #[cfg(fox_stub)] + { + println!(" llama.cpp: stub build (no native backend)"); + println!(" Flash Attn: not available"); + println!(" CUDA: not available"); + println!(" Metal: not available"); + } + + #[cfg(not(fox_stub))] + { + let has_cuda = + cfg!(any(target_os = "linux", target_os = "windows",)) && cuda_backend_present(); + + let has_metal = cfg!(target_os = "macos"); + + if has_cuda { + println!(" llama.cpp: built with CUDA support"); + } else if has_metal { + println!(" llama.cpp: built with Metal support"); + } else { + println!(" llama.cpp: CPU only"); + } + + println!( + " Flash Attn: {}", + if has_cuda || has_metal { + "available" + } else { + "not available" + } + ); + println!( + " CUDA: {}", + if has_cuda { + "available" + } else { + "not available" + } + ); + println!( + " Metal: {}", + if has_metal { + "available" + } else { + "not available" + } + ); + } +} + +#[cfg(not(fox_stub))] +fn cuda_backend_present() -> bool { + let exe = match std::env::current_exe() { + Ok(p) => p, + Err(_) => return false, + }; + let dir = match exe.parent() { + Some(d) => d, + None => return false, + }; + let so_name = if cfg!(target_os = "windows") { + "ggml-cuda.dll" + } else { + "libggml-cuda.so" + }; + dir.join(so_name).exists() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn gpu_info_runs_without_panic() { + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + let result = run_gpu_info(GpuInfoArgs {}).await; + assert!(result.is_ok()); + }); + } +} diff --git a/src/cli/mod.rs b/src/cli/mod.rs index 61f68c6..b99de24 100644 --- a/src/cli/mod.rs +++ b/src/cli/mod.rs @@ -5,6 +5,8 @@ pub mod bench; pub mod bench_kv; pub mod bench_prefill; pub mod bench_spec; +pub mod discover; +pub mod gpu_info; pub mod list; pub mod mcp; pub mod models; @@ -70,6 +72,10 @@ pub enum Command { Alias(alias::AliasArgs), /// Start an MCP (Model Context Protocol) server over stdio for IDE integration Mcp(mcp::McpArgs), + /// Scan well-known directories (HuggingFace, Ollama, LM Studio, custom) for GGUF models + Discover(discover::DiscoverArgs), + /// Display GPU backend, VRAM, and driver details + GpuInfo(gpu_info::GpuInfoArgs), } /// Known subcommand names — anything else is treated as `fox run `. @@ -90,6 +96,8 @@ const SUBCOMMANDS: &[&str] = &[ "search", "alias", "mcp", + "discover", + "gpu-info", "help", ]; @@ -128,5 +136,7 @@ pub async fn run() -> anyhow::Result<()> { Command::Search(args) => search::run_search(args).await, Command::Alias(args) => alias::run_alias(args).await, Command::Mcp(args) => mcp::run_mcp(args).await, + Command::Discover(args) => discover::run_discover(args).await, + Command::GpuInfo(args) => gpu_info::run_gpu_info(args).await, } } diff --git a/src/lib.rs b/src/lib.rs index 01041f5..ba16339 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -7,6 +7,7 @@ pub(crate) mod engine; pub(crate) mod kv_cache; pub(crate) mod mcp; pub mod metrics; +pub mod model_discovery; pub mod model_registry; pub mod registry; pub(crate) mod scheduler; diff --git a/src/model_discovery.rs b/src/model_discovery.rs new file mode 100644 index 0000000..dfe7974 --- /dev/null +++ b/src/model_discovery.rs @@ -0,0 +1,791 @@ +// Multi-path GGUF model discovery across well-known directories. + +use std::collections::{HashMap, HashSet}; +use std::fmt; +use std::fs::File; +use std::io::Read; +use std::path::{Path, PathBuf}; + +use walkdir::WalkDir; + +const GGUF_MAGIC: [u8; 4] = [0x47, 0x47, 0x55, 0x46]; // "GGUF" + +#[derive(Debug, Clone)] +pub struct DiscoveredModel { + pub name: String, + pub path: PathBuf, + pub size_bytes: u64, + pub source: ModelSource, + pub shard_paths: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ModelSource { + Ferrumox, + HuggingFace, + Ollama, + LmStudio, + Custom(PathBuf), +} + +impl fmt::Display for ModelSource { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Ferrumox => write!(f, "ferrumox"), + Self::HuggingFace => write!(f, "huggingface"), + Self::Ollama => write!(f, "ollama"), + Self::LmStudio => write!(f, "lmstudio"), + Self::Custom(_) => write!(f, "custom"), + } + } +} + +pub fn discover_models(extra_dirs: &[PathBuf]) -> Vec { + let home = match dirs::home_dir() { + Some(h) => h, + None => return vec![], + }; + + let mut raw: Vec = Vec::new(); + + // Ferrumox own model dir + let ferrumox_dir = dirs::cache_dir() + .unwrap_or_else(|| PathBuf::from(".")) + .join("ferrumox") + .join("models"); + scan_gguf_dir(&ferrumox_dir, 1, ModelSource::Ferrumox, &mut raw); + + // HuggingFace cache + let hf_dir = home.join(".cache/huggingface/hub"); + scan_gguf_dir(&hf_dir, 5, ModelSource::HuggingFace, &mut raw); + + // Ollama blobs + let ollama_dir = home.join(".ollama/models"); + scan_ollama(&ollama_dir, &mut raw); + + // LM Studio + let lmstudio_dir = home.join(".lmstudio/models"); + scan_gguf_dir(&lmstudio_dir, 4, ModelSource::LmStudio, &mut raw); + + // Custom dirs from --model-dirs + for dir in extra_dirs { + scan_gguf_dir(dir, 5, ModelSource::Custom(dir.clone()), &mut raw); + } + + // Group shards + let mut models = group_shards(raw); + + // Deduplicate by canonical path + dedup_by_path(&mut models); + + // Derive human-readable names (skip if already set, e.g. Ollama from manifests) + for m in &mut models { + if m.name.is_empty() { + m.name = derive_name(&m.path, &m.source); + } + } + + models.sort_by_key(|a| a.name.to_lowercase()); + + tracing::info!(count = models.len(), "model discovery complete"); + for m in &models { + tracing::info!( + name = %m.name, + source = %m.source, + size_bytes = m.size_bytes, + path = %m.path.display(), + "discovered model" + ); + } + + models +} + +fn scan_gguf_dir( + dir: &Path, + max_depth: usize, + source: ModelSource, + out: &mut Vec, +) { + if !dir.is_dir() { + return; + } + for entry in WalkDir::new(dir) + .max_depth(max_depth) + .follow_links(true) + .into_iter() + .filter_map(|e| e.ok()) + { + let path = entry.path(); + if !path.is_file() { + continue; + } + let ext = path.extension().and_then(|e| e.to_str()); + if ext.map(|e| e.eq_ignore_ascii_case("gguf")) != Some(true) { + continue; + } + let size = entry.metadata().map(|m| m.len()).unwrap_or(0); + out.push(DiscoveredModel { + name: String::new(), + path: path.to_path_buf(), + size_bytes: size, + source: source.clone(), + shard_paths: vec![], + }); + } +} + +fn scan_ollama(ollama_dir: &Path, out: &mut Vec) { + let blobs_dir = ollama_dir.join("blobs"); + if !blobs_dir.is_dir() { + return; + } + + // Build manifest name map: blob hash -> "namespace/model:tag" + let manifest_names = parse_ollama_manifests(ollama_dir); + + for entry in WalkDir::new(&blobs_dir) + .max_depth(1) + .into_iter() + .filter_map(|e| e.ok()) + { + let path = entry.path(); + if !path.is_file() { + continue; + } + if !has_gguf_magic(path) { + continue; + } + let size = entry.metadata().map(|m| m.len()).unwrap_or(0); + let fname = path + .file_name() + .and_then(|n| n.to_str()) + .unwrap_or_default(); + + // Try to match blob filename to a manifest entry. + // Ollama blob filenames look like "sha256-". + let name = manifest_names.get(fname).cloned().unwrap_or_default(); + + out.push(DiscoveredModel { + name, + path: path.to_path_buf(), + size_bytes: size, + source: ModelSource::Ollama, + shard_paths: vec![], + }); + } +} + +fn parse_ollama_manifests(ollama_dir: &Path) -> HashMap { + let manifests_dir = ollama_dir.join("manifests"); + let mut map = HashMap::new(); + if !manifests_dir.is_dir() { + return map; + } + + // Manifest layout: manifests/registry.ollama.ai/// + for entry in WalkDir::new(&manifests_dir) + .max_depth(5) + .into_iter() + .filter_map(|e| e.ok()) + { + let path = entry.path(); + if !path.is_file() { + continue; + } + + let content = match std::fs::read_to_string(path) { + Ok(c) => c, + Err(_) => continue, + }; + + let json: serde_json::Value = match serde_json::from_str(&content) { + Ok(v) => v, + Err(_) => continue, + }; + + // Extract model name from directory structure + let rel = match path.strip_prefix(&manifests_dir) { + Ok(r) => r, + Err(_) => continue, + }; + let components: Vec<&str> = rel + .components() + .filter_map(|c| c.as_os_str().to_str()) + .collect(); + + // Expected: ["registry.ollama.ai", namespace, model, tag] + let friendly_name = if components.len() >= 4 { + format!( + "{}:{}", + if components[1] == "library" { + components[2].to_string() + } else { + format!("{}/{}", components[1], components[2]) + }, + components[3] + ) + } else { + continue; + }; + + // Parse layers to find the model blob + if let Some(layers) = json.get("layers").and_then(|l| l.as_array()) { + for layer in layers { + let media_type = layer + .get("mediaType") + .and_then(|m| m.as_str()) + .unwrap_or_default(); + if media_type.contains("model") { + if let Some(digest) = layer.get("digest").and_then(|d| d.as_str()) { + // Digest is "sha256:", blob filename is "sha256-" + let blob_name = digest.replace(':', "-"); + map.insert(blob_name, friendly_name.clone()); + } + } + } + } + } + + map +} + +fn has_gguf_magic(path: &Path) -> bool { + let mut f = match File::open(path) { + Ok(f) => f, + Err(_) => return false, + }; + let mut buf = [0u8; 4]; + match f.read_exact(&mut buf) { + Ok(()) => buf == GGUF_MAGIC, + Err(_) => false, + } +} + +/// Detect shard patterns like `model-00001-of-00005.gguf` and group them. +fn group_shards(models: Vec) -> Vec { + let re_shard = regex_shard_pattern(); + let mut shard_groups: HashMap<(String, PathBuf), Vec> = HashMap::new(); + let mut standalone: Vec = Vec::new(); + + for m in models { + if let Some(fname) = m.path.file_name().and_then(|n| n.to_str()) { + if let Some(caps) = re_shard.captures(fname) { + let prefix = caps.name("prefix").unwrap().as_str().to_string(); + let parent = m.path.parent().unwrap_or(Path::new("")).to_path_buf(); + let key = (prefix, parent); + shard_groups.entry(key).or_default().push(m); + continue; + } + } + standalone.push(m); + } + + for (_, mut shards) in shard_groups { + shards.sort_by(|a, b| a.path.cmp(&b.path)); + let combined_size: u64 = shards.iter().map(|s| s.size_bytes).sum(); + let first = shards.first().unwrap(); + let shard_paths: Vec = shards.iter().map(|s| s.path.clone()).collect(); + standalone.push(DiscoveredModel { + name: first.name.clone(), + path: first.path.clone(), + size_bytes: combined_size, + source: first.source.clone(), + shard_paths, + }); + } + + standalone +} + +struct ShardRegex; + +struct ShardCaptures<'a> { + prefix: &'a str, +} + +struct ShardMatch<'a> { + text: &'a str, +} + +impl<'a> ShardMatch<'a> { + fn as_str(&self) -> &'a str { + self.text + } +} + +impl<'a> ShardCaptures<'a> { + fn name(&self, group: &str) -> Option> { + if group == "prefix" { + Some(ShardMatch { text: self.prefix }) + } else { + None + } + } +} + +impl ShardRegex { + fn captures<'a>(&self, text: &'a str) -> Option> { + // Match pattern: -NNNNN-of-NNNNN.gguf + let text_lower = text.to_lowercase(); + if !text_lower.ends_with(".gguf") { + return None; + } + let without_ext = &text[..text.len() - 5]; // strip .gguf + // Look for -NNNNN-of-NNNNN at the end + // The digits section can vary in length but must be consistent + let parts: Vec<&str> = without_ext.rsplitn(4, '-').collect(); + // rsplitn with limit 4 on "prefix-00001-of-00005" gives: + // ["00005", "of", "00001", "prefix"] + if parts.len() < 4 { + return None; + } + let total = parts[0]; + let of_part = parts[1]; + let index = parts[2]; + let prefix_text = parts[3]; + + if of_part != "of" { + return None; + } + if total.is_empty() + || index.is_empty() + || !total.chars().all(|c| c.is_ascii_digit()) + || !index.chars().all(|c| c.is_ascii_digit()) + { + return None; + } + + Some(ShardCaptures { + prefix: prefix_text, + }) + } +} + +fn regex_shard_pattern() -> ShardRegex { + ShardRegex +} + +fn dedup_by_path(models: &mut Vec) { + let mut seen = HashSet::new(); + models.retain(|m| { + let canonical = m.path.canonicalize().unwrap_or_else(|_| m.path.clone()); + seen.insert(canonical) + }); +} + +fn derive_name(path: &Path, source: &ModelSource) -> String { + match source { + ModelSource::HuggingFace => derive_hf_name(path), + ModelSource::LmStudio => derive_lmstudio_name(path), + _ => path + .file_stem() + .and_then(|s| s.to_str()) + .unwrap_or("unknown") + .to_string(), + } +} + +fn derive_hf_name(path: &Path) -> String { + // Look for "models--org--name" in the path + for ancestor in path.ancestors() { + let Some(name) = ancestor.file_name().and_then(|n| n.to_str()) else { + continue; + }; + let Some(rest) = name.strip_prefix("models--") else { + continue; + }; + let parts: Vec<&str> = rest.splitn(2, "--").collect(); + if parts.len() == 2 { + let model_name = format!("{}/{}", parts[0], parts[1]); + if let Some(stem) = path.file_stem().and_then(|s| s.to_str()) { + return format!("{}/{}", model_name, stem); + } + return model_name; + } + } + path.file_stem() + .and_then(|s| s.to_str()) + .unwrap_or("unknown") + .to_string() +} + +fn derive_lmstudio_name(path: &Path) -> String { + // LM Studio: ~/.lmstudio/models/{org}/{model-file}.gguf + // Extract the last two components before the filename + let components: Vec<&str> = path + .components() + .filter_map(|c| c.as_os_str().to_str()) + .collect(); + + // Find "models" in the path and take next component as org + for (i, c) in components.iter().enumerate() { + if *c == "models" && i + 2 < components.len() { + let org = components[i + 1]; + let stem = path + .file_stem() + .and_then(|s| s.to_str()) + .unwrap_or("unknown"); + return format!("{}/{}", org, stem); + } + } + + path.file_stem() + .and_then(|s| s.to_str()) + .unwrap_or("unknown") + .to_string() +} + +/// Parse semicolon-separated directory list into paths. +pub fn parse_model_dirs(input: &str) -> Vec { + input + .split(';') + .map(|s| s.trim()) + .filter(|s| !s.is_empty()) + .map(|s| { + let p = PathBuf::from(s); + crate::cli::expand_tilde(&p) + }) + .collect() +} + +/// Look up a model name in the discovered models list. +/// Returns the path of the first match (exact, then prefix, then contains). +pub fn resolve_discovered(name: &str, models: &[DiscoveredModel]) -> Option<(String, PathBuf)> { + let lower = name.to_lowercase().replace(':', "-"); + + // Exact match on name + for m in models { + if m.name.eq_ignore_ascii_case(name) || m.name.to_lowercase().replace(':', "-") == lower { + return Some((m.name.clone(), m.path.clone())); + } + } + + // Name ends with the query (e.g. "llama3" matches "meta-llama/llama3") + for m in models { + let m_lower = m.name.to_lowercase().replace(':', "-"); + if m_lower.ends_with(&lower) { + return Some((m.name.clone(), m.path.clone())); + } + } + + // Contains + for m in models { + let m_lower = m.name.to_lowercase().replace(':', "-"); + if m_lower.contains(&lower) { + return Some((m.name.clone(), m.path.clone())); + } + } + + // Also try matching against the file stem + for m in models { + if let Some(stem) = m.path.file_stem().and_then(|s| s.to_str()) { + if stem.to_lowercase().contains(&lower) { + return Some((m.name.clone(), m.path.clone())); + } + } + } + + None +} + +#[cfg(test)] +mod tests { + use super::*; + + // ── Name derivation ───────────────────────────────────────────────── + + #[test] + fn test_derive_hf_name() { + let path = PathBuf::from( + "/home/user/.cache/huggingface/hub/models--meta-llama--Meta-Llama-3.1-8B/snapshots/abc123/model.gguf", + ); + let name = derive_hf_name(&path); + assert_eq!(name, "meta-llama/Meta-Llama-3.1-8B/model"); + } + + #[test] + fn test_derive_hf_name_no_pattern() { + let path = PathBuf::from("/some/random/path/model.gguf"); + let name = derive_hf_name(&path); + assert_eq!(name, "model"); + } + + #[test] + fn test_derive_lmstudio_name() { + let path = PathBuf::from( + "/home/user/.lmstudio/models/TheBloke/Mistral-7B-Instruct-v0.2-GGUF/mistral-7b.gguf", + ); + let name = derive_lmstudio_name(&path); + assert_eq!(name, "TheBloke/mistral-7b"); + } + + #[test] + fn test_derive_ferrumox_name() { + let path = PathBuf::from("/home/user/.cache/ferrumox/models/llama-3.2-3b.gguf"); + let name = derive_name(&path, &ModelSource::Ferrumox); + assert_eq!(name, "llama-3.2-3b"); + } + + #[test] + fn test_derive_custom_name() { + let path = PathBuf::from("/data/models/my-model.gguf"); + let name = derive_name(&path, &ModelSource::Custom(PathBuf::from("/data/models"))); + assert_eq!(name, "my-model"); + } + + // ── Shard grouping ────────────────────────────────────────────────── + + #[test] + fn test_shard_detection() { + let re = regex_shard_pattern(); + let caps = re.captures("model-00001-of-00005.gguf"); + assert!(caps.is_some()); + assert_eq!(caps.unwrap().name("prefix").unwrap().as_str(), "model"); + } + + #[test] + fn test_shard_detection_no_match() { + let re = regex_shard_pattern(); + assert!(re.captures("model.gguf").is_none()); + assert!(re.captures("model-q4.gguf").is_none()); + } + + #[test] + fn test_group_shards_combines() { + let models = vec![ + DiscoveredModel { + name: String::new(), + path: PathBuf::from("/models/llama-00001-of-00003.gguf"), + size_bytes: 1000, + source: ModelSource::Ferrumox, + shard_paths: vec![], + }, + DiscoveredModel { + name: String::new(), + path: PathBuf::from("/models/llama-00002-of-00003.gguf"), + size_bytes: 1000, + source: ModelSource::Ferrumox, + shard_paths: vec![], + }, + DiscoveredModel { + name: String::new(), + path: PathBuf::from("/models/llama-00003-of-00003.gguf"), + size_bytes: 1000, + source: ModelSource::Ferrumox, + shard_paths: vec![], + }, + DiscoveredModel { + name: String::new(), + path: PathBuf::from("/models/standalone.gguf"), + size_bytes: 500, + source: ModelSource::Ferrumox, + shard_paths: vec![], + }, + ]; + + let grouped = group_shards(models); + assert_eq!(grouped.len(), 2); + + let sharded = grouped.iter().find(|m| !m.shard_paths.is_empty()).unwrap(); + assert_eq!(sharded.size_bytes, 3000); + assert_eq!(sharded.shard_paths.len(), 3); + + let single = grouped.iter().find(|m| m.shard_paths.is_empty()).unwrap(); + assert_eq!(single.size_bytes, 500); + } + + // ── Deduplication ─────────────────────────────────────────────────── + + #[test] + fn test_dedup_by_path() { + let tmp = tempfile::tempdir().unwrap(); + let file = tmp.path().join("model.gguf"); + std::fs::write(&file, b"test").unwrap(); + + let mut models = vec![ + DiscoveredModel { + name: "a".to_string(), + path: file.clone(), + size_bytes: 4, + source: ModelSource::Ferrumox, + shard_paths: vec![], + }, + DiscoveredModel { + name: "b".to_string(), + path: file.clone(), + size_bytes: 4, + source: ModelSource::Custom(tmp.path().to_path_buf()), + shard_paths: vec![], + }, + ]; + + dedup_by_path(&mut models); + assert_eq!(models.len(), 1); + } + + // ── parse_model_dirs ──────────────────────────────────────────────── + + #[test] + fn test_parse_model_dirs_semicolon() { + let dirs = parse_model_dirs("/path/one;/path/two"); + assert_eq!(dirs.len(), 2); + assert_eq!(dirs[0], PathBuf::from("/path/one")); + assert_eq!(dirs[1], PathBuf::from("/path/two")); + } + + #[test] + fn test_parse_model_dirs_empty() { + let dirs = parse_model_dirs(""); + assert!(dirs.is_empty()); + } + + #[test] + fn test_parse_model_dirs_single() { + let dirs = parse_model_dirs("/single/path"); + assert_eq!(dirs.len(), 1); + } + + // ── resolve_discovered ────────────────────────────────────────────── + + #[test] + fn test_resolve_discovered_exact() { + let models = vec![DiscoveredModel { + name: "qwen2.5:7b".to_string(), + path: PathBuf::from("/test/model.gguf"), + size_bytes: 100, + source: ModelSource::Ollama, + shard_paths: vec![], + }]; + let result = resolve_discovered("qwen2.5:7b", &models); + assert!(result.is_some()); + assert_eq!(result.unwrap().0, "qwen2.5:7b"); + } + + #[test] + fn test_resolve_discovered_contains() { + let models = vec![DiscoveredModel { + name: "meta-llama/Meta-Llama-3.1-8B/model-q4".to_string(), + path: PathBuf::from("/test/model.gguf"), + size_bytes: 100, + source: ModelSource::HuggingFace, + shard_paths: vec![], + }]; + let result = resolve_discovered("llama-3.1", &models); + assert!(result.is_some()); + } + + #[test] + fn test_resolve_discovered_not_found() { + let models = vec![DiscoveredModel { + name: "llama".to_string(), + path: PathBuf::from("/test/model.gguf"), + size_bytes: 100, + source: ModelSource::Ferrumox, + shard_paths: vec![], + }]; + let result = resolve_discovered("mistral", &models); + assert!(result.is_none()); + } + + // ── has_gguf_magic ────────────────────────────────────────────────── + + #[test] + fn test_has_gguf_magic_positive() { + let tmp = tempfile::tempdir().unwrap(); + let file = tmp.path().join("blob"); + std::fs::write(&file, b"GGUF\x03\x00\x00\x00").unwrap(); + assert!(has_gguf_magic(&file)); + } + + #[test] + fn test_has_gguf_magic_negative() { + let tmp = tempfile::tempdir().unwrap(); + let file = tmp.path().join("blob"); + std::fs::write(&file, b"NOT_GGUF").unwrap(); + assert!(!has_gguf_magic(&file)); + } + + #[test] + fn test_has_gguf_magic_too_short() { + let tmp = tempfile::tempdir().unwrap(); + let file = tmp.path().join("blob"); + std::fs::write(&file, b"GG").unwrap(); + assert!(!has_gguf_magic(&file)); + } + + #[test] + fn test_has_gguf_magic_missing_file() { + assert!(!has_gguf_magic(Path::new("/nonexistent/file"))); + } + + // ── ModelSource Display ───────────────────────────────────────────── + + #[test] + fn test_model_source_display() { + assert_eq!(format!("{}", ModelSource::Ferrumox), "ferrumox"); + assert_eq!(format!("{}", ModelSource::HuggingFace), "huggingface"); + assert_eq!(format!("{}", ModelSource::Ollama), "ollama"); + assert_eq!(format!("{}", ModelSource::LmStudio), "lmstudio"); + assert_eq!( + format!("{}", ModelSource::Custom(PathBuf::from("/x"))), + "custom" + ); + } + + // ── scan_gguf_dir ─────────────────────────────────────────────────── + + #[test] + fn test_scan_gguf_dir_finds_models() { + let tmp = tempfile::tempdir().unwrap(); + std::fs::write(tmp.path().join("model-a.gguf"), b"fake").unwrap(); + std::fs::write(tmp.path().join("model-b.gguf"), b"fake2").unwrap(); + std::fs::write(tmp.path().join("readme.txt"), b"ignore").unwrap(); + + let mut out = Vec::new(); + scan_gguf_dir(tmp.path(), 1, ModelSource::Ferrumox, &mut out); + assert_eq!(out.len(), 2); + } + + #[test] + fn test_scan_gguf_dir_nonexistent() { + let mut out = Vec::new(); + scan_gguf_dir( + Path::new("/does/not/exist"), + 1, + ModelSource::Ferrumox, + &mut out, + ); + assert!(out.is_empty()); + } + + // ── Ollama name derivation ────────────────────────────────────────── + + #[test] + fn test_ollama_name_from_manifest_library() { + let tmp = tempfile::tempdir().unwrap(); + let manifest_dir = tmp + .path() + .join("manifests/registry.ollama.ai/library/llama3/latest"); + std::fs::create_dir_all(manifest_dir.parent().unwrap()).unwrap(); + + let manifest = serde_json::json!({ + "layers": [{ + "mediaType": "application/vnd.ollama.image.model", + "digest": "sha256:abc123", + "size": 1000 + }] + }); + std::fs::write(&manifest_dir, manifest.to_string()).unwrap(); + + let blobs_dir = tmp.path().join("blobs"); + std::fs::create_dir_all(&blobs_dir).unwrap(); + std::fs::write(blobs_dir.join("sha256-abc123"), b"GGUF\x03\x00\x00\x00").unwrap(); + + let mut out = Vec::new(); + scan_ollama(tmp.path(), &mut out); + + assert_eq!(out.len(), 1); + assert_eq!(out[0].name, "llama3:latest"); + } +} From e0176cc4bb09171ce10df4598f490e479b8b9428 Mon Sep 17 00:00:00 2001 From: CrimsonMartin Date: Fri, 28 Aug 2026 21:14:35 +0000 Subject: [PATCH 3/6] ci: port secret scanning, llama.cpp auto-update, ROCm and .deb release lanes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Re-applies the fork's CI additions on the upstream base: - secret-scanning.yml: TruffleHog verified-secrets scan on push/PR. - update-llamacpp.yml: weekly job that bumps vendor/llama.cpp to upstream master, verifies a stub cargo check, and opens a labeled PR. - docker.yml: matrix build adding a -rocm image variant (Dockerfile.rocm) next to the existing CUDA image, with per-variant build cache and a disk-space step for the large ROCm base image. - release.yml: x86_64-linux-rocm release lane (ROCm 6.2 HIP SDK from AMD's repo, hipcc auto-detected by build.rs), plus a .deb package on the CPU leg via cargo-deb — smoke-tested by extracting and running the packaged binary, mirroring the existing tarball gate. The .deb asset globs lib*.so* so it cannot repeat the libmtmd.so.0 omission. --- .github/workflows/docker.yml | 29 +++++++- .github/workflows/release.yml | 63 +++++++++++++++++ .github/workflows/secret-scanning.yml | 24 +++++++ .github/workflows/update-llamacpp.yml | 98 +++++++++++++++++++++++++++ Cargo.toml | 21 ++++++ 5 files changed, 232 insertions(+), 3 deletions(-) create mode 100644 .github/workflows/secret-scanning.yml create mode 100644 .github/workflows/update-llamacpp.yml diff --git a/.github/workflows/docker.yml b/.github/workflows/docker.yml index 5e5ea58..64cb933 100644 --- a/.github/workflows/docker.yml +++ b/.github/workflows/docker.yml @@ -11,18 +11,40 @@ env: jobs: build: - name: Build & push (linux/amd64) + name: Build & push (${{ matrix.variant }}, linux/amd64) runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + include: + - variant: cuda + dockerfile: Dockerfile + suffix: "" + - variant: rocm + dockerfile: Dockerfile.rocm + suffix: "-rocm" steps: - uses: actions/checkout@v4 with: submodules: recursive + # The ROCm devel base image plus the build layers overflow the runner's + # default free space; clear the preinstalled toolchains we don't use. + - name: Free disk space + if: matrix.variant == 'rocm' + run: | + sudo rm -rf /usr/local/lib/android /usr/share/dotnet /opt/ghc + sudo apt-get remove -y '^dotnet-.*' '^llvm-.*' google-chrome-stable || true + sudo apt-get autoremove -y + df -h / + - name: Docker metadata (tags) id: meta uses: docker/metadata-action@v5 with: images: ${{ env.IMAGE }} + flavor: | + suffix=${{ matrix.suffix }},onlatest=true tags: | type=semver,pattern={{version}} type=semver,pattern={{major}}.{{minor}} @@ -42,9 +64,10 @@ jobs: uses: docker/build-push-action@v6 with: context: . + file: ${{ matrix.dockerfile }} platforms: linux/amd64 push: true tags: ${{ steps.meta.outputs.tags }} labels: ${{ steps.meta.outputs.labels }} - cache-from: type=gha - cache-to: type=gha,mode=max + cache-from: type=gha,scope=${{ matrix.variant }} + cache-to: type=gha,mode=max,scope=${{ matrix.variant }} diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 5279821..ce36f8e 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -74,12 +74,20 @@ jobs: runner: ubuntu-22.04 suffix: "" gpu_deps: "" + build_deb: true # Vulkan — needs Ubuntu 24.04 for glslc/spirv-headers. The resulting binary # needs glibc 2.39+ and a Vulkan driver at runtime (Mesa RADV/ANV, etc.). - name: x86_64-linux-vulkan runner: ubuntu-24.04 suffix: "-vulkan" gpu_deps: "glslc glslang-tools libvulkan-dev spirv-headers" + # ROCm — HIP SDK from AMD's own repo (the ROCm install step below); + # build.rs auto-enables the HIP backend when hipcc is on PATH. + - name: x86_64-linux-rocm + runner: ubuntu-22.04 + suffix: "-rocm" + gpu_deps: "" + gpu_backend: rocm steps: - uses: actions/checkout@v4 @@ -109,6 +117,27 @@ jobs: sudo apt-get update -q sudo apt-get install -y build-essential cmake clang libclang-dev ninja-build ${{ matrix.gpu_deps }} + - name: Install ROCm HIP SDK + if: matrix.gpu_backend == 'rocm' + run: | + sudo mkdir -p /etc/apt/keyrings + wget -q -O - https://repo.radeon.com/rocm/rocm.gpg.key \ + | sudo gpg --dearmor -o /etc/apt/keyrings/rocm.gpg + echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] \ + https://repo.radeon.com/rocm/apt/6.2 jammy main" \ + | sudo tee /etc/apt/sources.list.d/rocm.list + # Give ROCm repo higher priority than Ubuntu's packages. + # Ubuntu 22.04 ships rocminfo 5.0.0-1 (upstream versioning) which + # apt prefers over ROCm's 1.0.0.6XXXX (Radeon versioning). + # Pin-Priority 1001 forces ROCm packages even when they look like a downgrade. + printf 'Package: *\nPin: origin repo.radeon.com\nPin-Priority: 1001\n' \ + | sudo tee /etc/apt/preferences.d/rocm.pref + sudo apt-get update -q + sudo apt-get install -y rocm-hip-sdk + echo "/opt/rocm/bin" >> "$GITHUB_PATH" + echo "/opt/rocm/lib/llvm/bin" >> "$GITHUB_PATH" + echo "HIPCC=/opt/rocm/bin/hipcc" >> "$GITHUB_ENV" + - name: Build run: cargo build --release --bin fox --bin fox-bench @@ -133,12 +162,42 @@ jobs: echo "ASSET=$TARBALL" >> "$GITHUB_ENV" echo "ASSET_SHA256=${TARBALL}.sha256" >> "$GITHUB_ENV" + - name: Install cargo-deb + if: matrix.build_deb + run: cargo install cargo-deb --locked + + - name: Build .deb package + if: matrix.build_deb + shell: bash + run: | + cargo deb --no-build + DEB=$(ls target/debian/*.deb) + sha256sum "$DEB" > "${DEB}.sha256" + echo "DEB_ASSET=$DEB" >> "$GITHUB_ENV" + echo "DEB_ASSET_SHA256=${DEB}.sha256" >> "$GITHUB_ENV" + + - name: The .deb must actually install its binary + # Same reasoning as the tarball smoke test below: unpack what is about to + # be published and run it, instead of trusting the asset glob. + if: matrix.build_deb + shell: bash + run: | + SMOKE="$(mktemp -d)" + dpkg-deb -x "$DEB_ASSET" "$SMOKE" + LD_LIBRARY_PATH="$SMOKE/usr/lib" "$SMOKE/usr/bin/fox" --version + echo "packaged .deb installs a binary that starts and reports its version" + - name: The tarball must actually run # Unpack what is about to be published, somewhere else, and start it. v0.20.3 # shipped a binary that died on a missing `libmtmd.so.0` and nothing noticed, # because packaging was only ever checked by reading the glob. A release that # cannot print its own version is not a release. shell: bash + # The ROCm leg's libggml-hip.so pulls in the HIP runtime from /opt/rocm, + # which the runner has but the default loader path does not; empty and + # harmless on the other legs. + env: + LD_LIBRARY_PATH: ${{ matrix.gpu_backend == 'rocm' && '/opt/rocm/lib' || '' }} run: | SMOKE="$(mktemp -d)" tar xzf "$ASSET" -C "$SMOKE" @@ -155,9 +214,13 @@ jobs: # curl -fsSL .../releases/latest/download/install.sh | sh # and that URL 404'd, because the only assets uploaded were the tarball and # its checksum. The advertised way to install fox did not work. + # DEB_ASSET is only set on the build_deb leg; softprops skips the + # resulting empty lines on the others. files: | ${{ env.ASSET }} ${{ env.ASSET_SHA256 }} + ${{ env.DEB_ASSET }} + ${{ env.DEB_ASSET_SHA256 }} install.sh install.ps1 diff --git a/.github/workflows/secret-scanning.yml b/.github/workflows/secret-scanning.yml new file mode 100644 index 0000000..4d28bd3 --- /dev/null +++ b/.github/workflows/secret-scanning.yml @@ -0,0 +1,24 @@ +name: Secret Scanning + +on: + push: + branches: [main, develop] + pull_request: + branches: [main, develop] + +env: + FORCE_JAVASCRIPT_ACTIONS_TO_NODE24: true + +jobs: + secret-scan: + name: Secret Scan + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v5 + with: + fetch-depth: 0 + + - name: TruffleHog secret scan + uses: trufflesecurity/trufflehog@v3.88.26 + with: + extra_args: --only-verified diff --git a/.github/workflows/update-llamacpp.yml b/.github/workflows/update-llamacpp.yml new file mode 100644 index 0000000..aab109f --- /dev/null +++ b/.github/workflows/update-llamacpp.yml @@ -0,0 +1,98 @@ +name: Update llama.cpp + +on: + schedule: + - cron: '0 6 * * 1' # Monday 6 AM UTC + workflow_dispatch: + +permissions: + contents: write + pull-requests: write + issues: write + +env: + SUBMODULE_PATH: vendor/llama.cpp + +jobs: + update: + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v5 + with: + token: ${{ secrets.PAT_TOKEN }} + submodules: recursive + fetch-depth: 0 + + - name: Check for upstream updates + id: check + working-directory: ${{ env.SUBMODULE_PATH }} + run: | + git fetch origin master + CURRENT=$(git rev-parse HEAD) + LATEST=$(git rev-parse origin/master) + BEHIND=$(git rev-list --count HEAD..origin/master) + echo "current=$CURRENT" >> "$GITHUB_OUTPUT" + echo "latest=$LATEST" >> "$GITHUB_OUTPUT" + echo "behind=$BEHIND" >> "$GITHUB_OUTPUT" + if [ "$CURRENT" = "$LATEST" ]; then + echo "Already up to date with llama.cpp master" + else + echo "$BEHIND new commits available" + fi + + - name: Update submodule + if: steps.check.outputs.behind != '0' + working-directory: ${{ env.SUBMODULE_PATH }} + run: git checkout origin/master + + - name: Verify build + if: steps.check.outputs.behind != '0' + run: | + FOX_SKIP_LLAMA=1 rustup run stable cargo check 2>&1 + env: + CARGO_TERM_COLOR: always + + - name: Create pull request + if: steps.check.outputs.behind != '0' + run: | + BEHIND=${{ steps.check.outputs.behind }} + SHORT_OLD=$(echo "${{ steps.check.outputs.current }}" | cut -c1-10) + SHORT_NEW=$(echo "${{ steps.check.outputs.latest }}" | cut -c1-10) + BRANCH="auto/llama-cpp-${SHORT_NEW}" + + git config user.name "github-actions[bot]" + git config user.email "github-actions[bot]@users.noreply.github.com" + + # Close any existing auto-update PRs + gh pr list --label "llama.cpp-update" --state open --json number --jq '.[].number' \ + | xargs -I{} gh pr close {} --comment "Superseded by newer update" 2>/dev/null || true + + git push origin --delete "$BRANCH" 2>/dev/null || true + git checkout -b "$BRANCH" + git add "${{ env.SUBMODULE_PATH }}" + git commit -m "$(cat < Date: Fri, 28 Aug 2026 21:39:48 +0000 Subject: [PATCH 4/6] =?UTF-8?q?feat:=20parallel=20vision=20preprocessing?= =?UTF-8?q?=20=E2=80=94=20mtmd=20context=20pool,=20CLIP=20cache,=20split?= =?UTF-8?q?=20prefill?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ports the fork's vision-parallelism work onto the current vision implementation, redesigned around today's mtmd helper API. Before this, a multimodal prefill was one atomic mtmd_helper_eval_chunks call that ran the CLIP encode while holding the llama context lock, so every other request — text ones included — stalled behind each image. Now tokenize_multimodal CLIP-encodes each media chunk up front (mtmd_encode_chunk + mtmd_get_output_embd) and stores the embeddings on MultimodalChunks; prefill walks the chunks and decodes text as tokens and images from the stored embeddings via mtmd_helper_decode_image_chunk, which handles M-RoPE and non-causal attention internally. The atomic path remains as fallback for chunks without embeddings. - MtmdPool: mtmd contexts become a checkout pool. --vision-contexts N (FOX_VISION_CONTEXTS, default 1) sizes it, so N requests encode in parallel; N=1 preserves today's serialization, but the encode still moves off the engine thread. - CLIP cache: per-model LRU (16 images) keyed by the raw image bytes' hash — chat clients re-send the history's images every turn, and each now encodes once. - Handlers call prepare_multimodal_prompt through a spawn_blocking wrapper so the encode doesn't stall tokio workers. Validated: make ci (stub fmt/clippy/tests/build + real cargo check), real-build clippy clean, 521 real-mode lib tests pass, docs flag check. Runtime validation on GPU hardware with a vision model still pending — this changes prefill mechanics for every multimodal request. --- docs/cli/serve.md | 1 + src/api/auth.rs | 1 + src/api/ollama/chat.rs | 18 +-- src/api/ollama/generate.rs | 16 +-- src/api/ollama/management.rs | 1 + src/api/shared/inference.rs | 36 ++++++ src/api/test_helpers.rs | 4 + src/api/v1/chat.rs | 20 +-- src/api/v1/models.rs | 1 + src/cli/bench.rs | 1 + src/cli/bench_kv.rs | 1 + src/cli/bench_prefill.rs | 1 + src/cli/bench_spec.rs | 1 + src/cli/probe.rs | 1 + src/cli/run.rs | 1 + src/cli/serve.rs | 8 ++ src/engine/model/llama_cpp/batch.rs | 102 +++++++++++++--- src/engine/model/llama_cpp/golden.rs | 2 + src/engine/model/llama_cpp/mod.rs | 162 +++++++++++++++++++++---- src/engine/model/llama_cpp/stub.rs | 3 +- src/engine/model/llama_cpp/vocab.rs | 174 +++++++++++++++++++-------- src/engine/model/mod.rs | 35 ++++++ src/mcp/mod.rs | 20 +-- src/model_registry/config.rs | 6 + src/model_registry/loader.rs | 3 + src/model_registry/mod.rs | 1 + 26 files changed, 488 insertions(+), 132 deletions(-) diff --git a/docs/cli/serve.md b/docs/cli/serve.md index 506d875..ded0aad 100644 --- a/docs/cli/serve.md +++ b/docs/cli/serve.md @@ -67,6 +67,7 @@ fox serve --json-logs --port 8080 --max-models 2 --keep-alive-secs 600 | `--spec-draft-len ` | `FOX_SPEC_DRAFT_LEN` | `4` | Maximum draft tokens proposed per speculative step. | | `--draft-model ` | `FOX_DRAFT_MODEL` | — | Name/path of a smaller model to use as the speculative-decoding draft proposer instead of n-gram lookup — generalizes speculation to any text, not just repetitive output. Requires `--speculative true` (ignored with a startup warning otherwise). The draft and target must share the same tokenizer — checked at load time, fails loudly on mismatch. Loaded once alongside the target and kept resident for the process lifetime; not subject to LRU eviction or VRAM budgeting — size both models to fit. `--spec-ngram` is ignored in this mode. | | `--mmproj ` | `FOX_MMPROJ` | — | Name/path of a paired mmproj (vision projector) GGUF, enabling image input (OpenAI `image_url` / Ollama `images`) via llama.cpp's `mtmd` library. Resolved the same way as `--model-path`/`--draft-model` (alias, disk match, or a direct file path). One global pairing — matched against whatever model is currently loaded, like `--draft-model`. Only `data:` base64 image URIs are accepted (no remote fetch). See `docs/design/vision-support.md`. | +| `--vision-contexts ` | `FOX_VISION_CONTEXTS` | `1` | Number of vision (mtmd) contexts loaded with the mmproj. `1` = concurrent requests take turns CLIP-encoding their images; `N` = up to N encode in parallel, at roughly one mmproj's VRAM footprint per extra context. Repeated images are additionally served from a per-model CLIP cache, so a conversation that re-sends its history's images every turn only pays each encode once. Ignored without `--mmproj`. | | `--lora-modules =[:][,...]` | `FOX_LORA_MODULES` | — | Comma-separated LoRA adapters loaded onto the primary model (`--model-path`) at startup, e.g. `finetune=/path/adapter.gguf:0.8,other=/other.gguf`. `SCALE` defaults to `1.0`. Select an adapter per-request by passing its `NAME` as the `model` field in `/v1/chat/completions` or `/api/chat` — the request is served by the primary model with that adapter applied. Requests without a recognized adapter name use the primary model unmodified. Adapters are a property of the whole context, not per-sequence: concurrent requests on different adapters are grouped and processed as separate sub-batches (see `docs/design/lora-support.md`), and prefix caching is skipped for any request carrying an adapter selection. | | `--gpu-memory-fraction ` | `FOX_GPU_MEMORY_FRACTION` | `0.85` | Fraction of GPU VRAM reserved for the KV cache. Must be between 0.0 and 1.0. The remaining memory is left for model weights and other allocations. | | `--type-kv ` | `FOX_TYPE_KV` | `f16` | KV cache element type for both K and V: `f16`, `q8_0`, or `q4_0`. | diff --git a/src/api/auth.rs b/src/api/auth.rs index 1d7b992..c9cf447 100644 --- a/src/api/auth.rs +++ b/src/api/auth.rs @@ -83,6 +83,7 @@ mod tests { n_gpu_layers: -1, moe_offload_cpu: false, mmproj: None, + vision_contexts: 1, mtp_model: None, lora_modules: Vec::new(), primary_model: None, diff --git a/src/api/ollama/chat.rs b/src/api/ollama/chat.rs index 9a9f29d..3dc80fd 100644 --- a/src/api/ollama/chat.rs +++ b/src/api/ollama/chat.rs @@ -13,7 +13,7 @@ use base64::{engine::general_purpose::STANDARD, Engine as _}; use crate::api::error::load_model_or_respond; use crate::api::router::AppState; use crate::api::shared::inference::{ - extract_thinking, parse_tool_call, prepare_multimodal_prompt, prepare_prompt, + extract_thinking, parse_tool_call, prepare_multimodal_prompt_blocking, prepare_prompt, resolve_tool_call_parser, resolve_tool_choice, sampling_from_ollama, MessageForTemplate, }; use crate::api::shared::streaming::{ @@ -189,17 +189,19 @@ pub async fn ollama_chat( } let (prompt_tokens, prompt_tokens_len, multimodal) = if use_vision { - match prepare_multimodal_prompt( - &entry, + match prepare_multimodal_prompt_blocking( + entry.clone(), messages, - state.system_prompt.as_deref(), - eff_tools, + state.system_prompt.clone(), + eff_tools.map(<[_]>::to_vec), false, // tool_required (Ollama uses auto only) None, // specific_tool - response_format.as_ref(), + response_format.clone(), show_thinking_in_output, - &images, - ) { + images, + ) + .await + { Ok(chunks) => { let n = chunks.n_positions(); (Vec::new(), n, Some(chunks)) diff --git a/src/api/ollama/generate.rs b/src/api/ollama/generate.rs index eca35d5..810d533 100644 --- a/src/api/ollama/generate.rs +++ b/src/api/ollama/generate.rs @@ -12,7 +12,7 @@ use base64::{engine::general_purpose::STANDARD, Engine as _}; use crate::api::error::load_model_or_respond; use crate::api::router::AppState; use crate::api::shared::inference::{ - prepare_multimodal_prompt, prepare_prompt, sampling_from_ollama, MessageForTemplate, + prepare_multimodal_prompt_blocking, prepare_prompt, sampling_from_ollama, MessageForTemplate, }; use crate::api::shared::streaming::{ collect_tokens_timed, ndjson_response, ndjson_stream, now_rfc3339, ollama_done_reason, @@ -148,17 +148,19 @@ pub async fn ollama_generate( } let (prompt_tokens, prompt_tokens_len, multimodal) = if use_vision { - match prepare_multimodal_prompt( - &entry, + match prepare_multimodal_prompt_blocking( + entry.clone(), messages, - state.system_prompt.as_deref(), + state.system_prompt.clone(), None, // no tools on /api/generate false, None, - response_format.as_ref(), + response_format.clone(), false, // show_thinking always false for /api/generate - &images, - ) { + images, + ) + .await + { Ok(chunks) => { let n = chunks.n_positions(); (Vec::new(), n, Some(chunks)) diff --git a/src/api/ollama/management.rs b/src/api/ollama/management.rs index 6ac3c22..e4b6182 100644 --- a/src/api/ollama/management.rs +++ b/src/api/ollama/management.rs @@ -479,6 +479,7 @@ mod tests { n_gpu_layers: -1, moe_offload_cpu: false, mmproj: None, + vision_contexts: 1, mtp_model: None, lora_modules: Vec::new(), primary_model: None, diff --git a/src/api/shared/inference.rs b/src/api/shared/inference.rs index 1b0b72a..b2c9061 100644 --- a/src/api/shared/inference.rs +++ b/src/api/shared/inference.rs @@ -747,6 +747,42 @@ pub fn prepare_multimodal_prompt( .tokenize_multimodal(&flat, show_thinking, tools_json.as_ref(), images) } +/// [`prepare_multimodal_prompt`], moved off the async runtime. Tokenizing a +/// multimodal prompt now CLIP-encodes its images on a pooled mtmd context +/// (hundreds of milliseconds to seconds of work), so handlers call this +/// `spawn_blocking` wrapper instead of stalling a tokio worker thread — which +/// is also what lets concurrent requests actually use `--vision-contexts N` +/// contexts in parallel. Arguments are owned because the closure outlives the +/// caller's borrows. +#[allow(clippy::too_many_arguments)] +pub async fn prepare_multimodal_prompt_blocking( + entry: std::sync::Arc, + messages: Vec, + system_prompt: Option, + tools: Option>, + tool_required: bool, + specific_tool: Option, + response_format: Option, + show_thinking: bool, + images: Vec>, +) -> anyhow::Result { + tokio::task::spawn_blocking(move || { + prepare_multimodal_prompt( + &entry, + messages, + system_prompt.as_deref(), + tools.as_deref(), + tool_required, + specific_tool.as_deref(), + response_format.as_ref(), + show_thinking, + &images, + ) + }) + .await + .map_err(|e| anyhow::anyhow!("vision preprocessing task panicked: {e}"))? +} + // --------------------------------------------------------------------------- // Unit tests // --------------------------------------------------------------------------- diff --git a/src/api/test_helpers.rs b/src/api/test_helpers.rs index a46d1bf..8cb2eba 100644 --- a/src/api/test_helpers.rs +++ b/src/api/test_helpers.rs @@ -54,6 +54,7 @@ pub fn make_test_registry( n_gpu_layers: -1, moe_offload_cpu: false, mmproj: None, + vision_contexts: 1, mtp_model: None, lora_modules: Vec::new(), primary_model: None, @@ -124,6 +125,7 @@ pub fn make_test_state_with_queue_depth( n_gpu_layers: -1, moe_offload_cpu: false, mmproj: None, + vision_contexts: 1, mtp_model: None, lora_modules: Vec::new(), primary_model: None, @@ -183,6 +185,7 @@ pub fn make_test_state_speculative( n_gpu_layers: -1, moe_offload_cpu: false, mmproj: None, + vision_contexts: 1, mtp_model: None, lora_modules: Vec::new(), primary_model: None, @@ -240,6 +243,7 @@ pub fn make_test_state_thinking(name: &str, dir: &std::path::Path) -> (AppState, n_gpu_layers: -1, moe_offload_cpu: false, mmproj: None, + vision_contexts: 1, mtp_model: None, lora_modules: Vec::new(), primary_model: None, diff --git a/src/api/v1/chat.rs b/src/api/v1/chat.rs index 9fd3d77..bd91042 100644 --- a/src/api/v1/chat.rs +++ b/src/api/v1/chat.rs @@ -14,7 +14,7 @@ use uuid::Uuid; use crate::api::error::{load_model_or_respond, AppError}; use crate::api::router::AppState; use crate::api::shared::inference::{ - parse_tool_call, prepare_multimodal_prompt, prepare_prompt, resolve_tool_call_parser, + parse_tool_call, prepare_multimodal_prompt_blocking, prepare_prompt, resolve_tool_call_parser, resolve_tool_choice, MessageForTemplate, }; use crate::api::shared::sampling_defaults as defaults; @@ -118,17 +118,19 @@ pub async fn chat_completions( tool_call_id: m.tool_call_id.clone(), }); } - match prepare_multimodal_prompt( - &entry, + match prepare_multimodal_prompt_blocking( + entry.clone(), messages, - state.system_prompt.as_deref(), - eff_tools, + state.system_prompt.clone(), + eff_tools.map(<[_]>::to_vec), tool_required, - specific_tool, - req.response_format.as_ref(), + specific_tool.map(str::to_string), + req.response_format.clone(), enable_thinking, - &images, - ) { + images, + ) + .await + { Ok(chunks) => { let n = chunks.n_positions(); (Vec::new(), n, Some(chunks)) diff --git a/src/api/v1/models.rs b/src/api/v1/models.rs index 4d8670c..c4d162b 100644 --- a/src/api/v1/models.rs +++ b/src/api/v1/models.rs @@ -226,6 +226,7 @@ mod tests { n_gpu_layers: -1, moe_offload_cpu: false, mmproj: None, + vision_contexts: 1, mtp_model: None, lora_modules: Vec::new(), primary_model: None, diff --git a/src/cli/bench.rs b/src/cli/bench.rs index 71504bd..946170c 100644 --- a/src/cli/bench.rs +++ b/src/cli/bench.rs @@ -122,6 +122,7 @@ pub async fn run_bench(args: BenchArgs) -> Result<()> { &tensor_split_parsed, args.moe_cpu, None, // mmproj_path + 1, // vision_contexts — no mmproj, single (unused) pool slot &[], // lora_modules false, // reranking — benches generate, never score, 0, // rs_rollback — no prompt reuse in this path diff --git a/src/cli/bench_kv.rs b/src/cli/bench_kv.rs index 1e7ea70..68d148e 100644 --- a/src/cli/bench_kv.rs +++ b/src/cli/bench_kv.rs @@ -326,6 +326,7 @@ pub async fn run_bench_kv(args: BenchKvArgs) -> Result<()> { &tensor_split, args.moe_cpu, None, // mmproj_path + 1, // vision_contexts — no mmproj, single (unused) pool slot &[], // lora_modules false, // reranking — benches generate, never score, 0, // rs_rollback — no prompt reuse in this path diff --git a/src/cli/bench_prefill.rs b/src/cli/bench_prefill.rs index d5fa2a4..cc54a5c 100644 --- a/src/cli/bench_prefill.rs +++ b/src/cli/bench_prefill.rs @@ -364,6 +364,7 @@ pub async fn run_bench_prefill(args: BenchPrefillArgs) -> Result<()> { &tensor_split, args.moe_cpu, None, // mmproj_path + 1, // vision_contexts — no mmproj, single (unused) pool slot &[], // lora_modules false, // reranking — benches generate, never score, 0, // rs_rollback — no prompt reuse in this path diff --git a/src/cli/bench_spec.rs b/src/cli/bench_spec.rs index a73f2e0..cb57493 100644 --- a/src/cli/bench_spec.rs +++ b/src/cli/bench_spec.rs @@ -261,6 +261,7 @@ pub async fn run_bench_spec(args: BenchSpecArgs) -> Result<()> { &tensor_split, args.moe_cpu, None, // mmproj_path + 1, // vision_contexts — no mmproj, single (unused) pool slot &[], // lora_modules false, // reranking — benches generate, never score, 0, // rs_rollback — no prompt reuse in this path diff --git a/src/cli/probe.rs b/src/cli/probe.rs index 630008f..21efb74 100644 --- a/src/cli/probe.rs +++ b/src/cli/probe.rs @@ -45,6 +45,7 @@ pub async fn run_probe(args: ProbeArgs) -> Result<()> { &[], // tensor_split false, // moe_offload_cpu None, // mmproj_path + 1, // vision_contexts — no mmproj, single (unused) pool slot &[], // lora_modules false, // reranking — benches generate, never score, 0, // rs_rollback — no prompt reuse in this path diff --git a/src/cli/run.rs b/src/cli/run.rs index 6f599cd..b7e3695 100644 --- a/src/cli/run.rs +++ b/src/cli/run.rs @@ -218,6 +218,7 @@ pub async fn run_run(args: RunArgs) -> Result<()> { &tensor_split_parsed, args.moe_cpu, None, // mmproj_path — `fox run` has no --mmproj flag yet; use `fox serve` for vision + 1, // vision_contexts — no mmproj, single (unused) pool slot &[], // lora_modules — same: fox run has no --lora-modules flag yet false, // reranking — benches generate, never score, 0, // rs_rollback — no prompt reuse in this path diff --git a/src/cli/serve.rs b/src/cli/serve.rs index 766cff0..1336727 100644 --- a/src/cli/serve.rs +++ b/src/cli/serve.rs @@ -165,6 +165,13 @@ pub struct ServeArgs { #[arg(long, env = "FOX_MMPROJ")] pub mmproj: Option, + /// Number of vision (mtmd) contexts loaded with the mmproj. 1 = concurrent + /// requests take turns CLIP-encoding their images; N = up to N encode in + /// parallel, at roughly one mmproj's VRAM footprint per extra context. + /// Ignored without --mmproj. + #[arg(long, default_value = "1", env = "FOX_VISION_CONTEXTS")] + pub vision_contexts: usize, + /// EXPERIMENTAL, opt-in, and BROKEN — do not enable it except to work on it. The /// head's `llama_decode` returns `-1` on every call, so it has never actually /// drafted: it returns a frozen, context-blind candidate set (asked to count to @@ -561,6 +568,7 @@ pub async fn run_serve(args: ServeArgs) -> Result<()> { spec_draft_len: args.spec_draft_len, draft_model: args.draft_model, mmproj: args.mmproj, + vision_contexts: args.vision_contexts, mtp_model: args.mtp_model, lora_modules: args .lora_modules diff --git a/src/engine/model/llama_cpp/batch.rs b/src/engine/model/llama_cpp/batch.rs index a88bbd9..da57fbb 100644 --- a/src/engine/model/llama_cpp/batch.rs +++ b/src/engine/model/llama_cpp/batch.rs @@ -491,8 +491,9 @@ impl LlamaCppModel { req_id: u64, req: &InferenceRequestForModel, ) -> Result { - let mtmd_ctx = self - .mtmd_ctx + let pool = self + .mtmd_pool + .as_ref() .ok_or_else(|| anyhow!("multimodal request but model has no mmproj loaded"))?; let chunks = req .multimodal @@ -505,25 +506,92 @@ impl LlamaCppModel { .map_err(|e| anyhow!("lock poisoned: {}", e))?; let ctx = ctx_guard.as_ptr(); + // The helpers still read mrope/non-causal config off an mtmd context, so + // one is checked out here too — but on the pre-encoded path below it never + // runs the CLIP encode, which is the expensive part that used to execute + // inside this function while every other request waited on `_ctx`. + let mtmd_nn = pool.acquire(); + let mtmd_ptr = mtmd_nn.as_ptr(); + // n_past is always 0: multimodal requests skip fox's chunked-prefill and // prefix-cache machinery entirely (see prompt_tokens being left empty for // them), so this is always a fresh sequence's first and only prefill call. let mut new_n_past: i32 = 0; - let ret = unsafe { - ffi::mtmd_helper_eval_chunks( - mtmd_ctx.as_ptr(), - ctx, - chunks.as_raw() as *const ffi::mtmd_input_chunks, - 0, - req.kv_seq_id.raw(), - self.effective_ctx as i32, - true, // logits_last - &mut new_n_past, - ) - }; - if ret != 0 { - return Err(anyhow!("mtmd_helper_eval_chunks failed (code {ret})")); - } + let eval_result = (|| -> Result<()> { + if chunks.has_image_embeddings() { + // Split path: tokenize time already CLIP-encoded every media chunk + // (`tokenize_multimodal_impl`), so walk the chunks and decode — + // text chunks as tokens, media chunks from their stored embeddings. + // Both helpers handle the model quirks (M-RoPE positions, Gemma's + // non-causal attention) that a manual llama_batch would have to + // re-implement. + let chunks_raw = chunks.as_raw() as *const ffi::mtmd_input_chunks; + let n_chunks = unsafe { ffi::mtmd_input_chunks_size(chunks_raw) }; + for i in 0..n_chunks { + let chunk = unsafe { ffi::mtmd_input_chunks_get(chunks_raw, i) }; + let is_last = i + 1 == n_chunks; + let ret = match chunks.embedding_for_chunk(i) { + Some(embd) => unsafe { + ffi::mtmd_helper_decode_image_chunk( + mtmd_ptr, + ctx, + chunk, + embd.as_ptr() as *mut f32, + new_n_past, + req.kv_seq_id.raw(), + self.effective_ctx as i32, + &mut new_n_past, + None, + std::ptr::null_mut(), + ) + }, + // Text chunks (and any media chunk that somehow has no + // stored embedding) go through the single-chunk eval, + // which encodes on demand for media and is a plain + // llama_decode for text. + None => unsafe { + ffi::mtmd_helper_eval_chunk_single( + mtmd_ptr, + ctx, + chunk, + new_n_past, + req.kv_seq_id.raw(), + self.effective_ctx as i32, + is_last, // logits_last + &mut new_n_past, + ) + }, + }; + if ret != 0 { + return Err(anyhow!( + "multimodal prefill failed on chunk {i} (code {ret})" + )); + } + } + Ok(()) + } else { + // Atomic fallback: no pre-encoded embeddings on this prompt, so + // mtmd runs encode + decode itself (the pre-0.23 behaviour). + let ret = unsafe { + ffi::mtmd_helper_eval_chunks( + mtmd_ptr, + ctx, + chunks.as_raw() as *const ffi::mtmd_input_chunks, + 0, + req.kv_seq_id.raw(), + self.effective_ctx as i32, + true, // logits_last + &mut new_n_past, + ) + }; + if ret != 0 { + return Err(anyhow!("mtmd_helper_eval_chunks failed (code {ret})")); + } + Ok(()) + } + })(); + pool.release(mtmd_nn); + eval_result?; let n_vocab = self.config.vocab_size as i32; let logits_ptr = unsafe { ffi::llama_get_logits_ith(ctx, -1) }; diff --git a/src/engine/model/llama_cpp/golden.rs b/src/engine/model/llama_cpp/golden.rs index b240453..b5397d3 100644 --- a/src/engine/model/llama_cpp/golden.rs +++ b/src/engine/model/llama_cpp/golden.rs @@ -41,6 +41,7 @@ fn golden_model() -> Option { &[], // tensor_split false, // moe_offload_cpu None, // mmproj_path + 1, // vision_contexts — no mmproj, single (unused) pool slot &[], // lora_modules false, // reranking — golden tests exercise generation, 0, // rs_rollback — no prompt reuse in this path @@ -673,6 +674,7 @@ fn golden_draft_model_speculative_matches_greedy() { &[], // tensor_split false, // moe_offload_cpu None, // mmproj_path + 1, // vision_contexts — no mmproj, single (unused) pool slot &[], // lora_modules false, // reranking — golden tests exercise generation, 0, // rs_rollback — no prompt reuse in this path diff --git a/src/engine/model/llama_cpp/mod.rs b/src/engine/model/llama_cpp/mod.rs index d9ea671..320379c 100644 --- a/src/engine/model/llama_cpp/mod.rs +++ b/src/engine/model/llama_cpp/mod.rs @@ -394,6 +394,74 @@ impl Drop for GrammarSampler { } } +/// One cached CLIP encode: keyed by (hash, byte length) of the raw image bytes, +/// holding the `n_image_tokens × n_embd_inp` floats `mtmd_get_output_embd` produced. +#[cfg(not(fox_stub))] +type ClipCache = lru::LruCache<(u64, usize), Arc>>; + +/// CLIP-cache capacity, in images. Entries are `n_image_tokens × n_embd_inp` +/// floats — roughly 1-4 MiB for typical vision models — so the cache tops out +/// around tens of MiB of host RAM. Sized for the common shape of the win: a +/// handful of conversations each re-sending the same image every turn. +#[cfg(not(fox_stub))] +const CLIP_CACHE_ENTRIES: usize = 16; + +/// Checkout pool of `mtmd_context`s for CLIP encoding. +/// +/// An `mtmd_context` carries per-call state (`mtmd_get_output_embd`'s buffer), so +/// one context serves one encode at a time; `acquire` blocks until a context is +/// free. With `--vision-contexts 1` (the default) this degenerates to a mutex +/// around the single context — exactly the exclusivity the previous +/// `Option>` field provided implicitly by only ever being +/// used under the engine's serialization. +#[cfg(not(fox_stub))] +pub(super) struct MtmdPool { + contexts: std::sync::Mutex>>, + available: std::sync::Condvar, + /// The model's input embedding width — sizes the per-image-token embedding + /// rows `mtmd_get_output_embd` returns. + pub(super) n_embd_inp: usize, +} + +#[cfg(not(fox_stub))] +impl MtmdPool { + fn new(contexts: Vec>, n_embd_inp: usize) -> Self { + Self { + contexts: std::sync::Mutex::new(contexts), + available: std::sync::Condvar::new(), + n_embd_inp, + } + } + + /// Check out a context for exclusive use; blocks while all are in flight. + /// Every `acquire` must be paired with a `release` on all exit paths — a + /// dropped checkout permanently shrinks the pool. + pub(super) fn acquire(&self) -> NonNull { + let mut pool = self.contexts.lock().expect("mtmd pool lock poisoned"); + loop { + if let Some(ctx) = pool.pop() { + return ctx; + } + pool = self.available.wait(pool).expect("mtmd pool lock poisoned"); + } + } + + pub(super) fn release(&self, ctx: NonNull) { + if let Ok(mut pool) = self.contexts.lock() { + pool.push(ctx); + } + self.available.notify_one(); + } + + fn free_all(&self) { + if let Ok(mut pool) = self.contexts.lock() { + for ctx in pool.drain(..) { + unsafe { ffi::mtmd_free(ctx.as_ptr()) }; + } + } + } +} + /// Llama.cpp model via FFI. #[cfg(not(fox_stub))] pub struct LlamaCppModel { @@ -430,10 +498,19 @@ pub struct LlamaCppModel { /// via `Model::bisection_retry_count()` and diffed into a Prometheus counter in /// `run_loop`, same pattern as `spec_proposed`/`spec_accepted`. pub(super) decode_bisection_retries: std::sync::atomic::AtomicU64, - /// Vision/multimodal context (`mtmd`), present only when this model was loaded + /// Vision/multimodal contexts (`mtmd`), present only when this model was loaded /// with a paired mmproj GGUF. `None` for the overwhelming majority of models — - /// every multimodal code path is gated on this being `Some`. - pub(super) mtmd_ctx: Option>, + /// every multimodal code path is gated on this being `Some`. Holds one context + /// by default; `--vision-contexts N` creates a pool of N so concurrent requests + /// CLIP-encode their images in parallel (each context is checked out exclusively, + /// so per-context state like `mtmd_get_output_embd`'s buffer is never shared). + pub(super) mtmd_pool: Option, + /// LRU of CLIP embeddings keyed by (hash, byte length) of the raw image bytes: + /// a client that re-sends the same image every conversation turn (which is how + /// both the OpenAI and Ollama chat APIs work — history carries the image) pays + /// the encode once. Entry values are shared into `MultimodalChunks` via `Arc`. + /// Always present; only the vision path ever touches it. + pub(super) clip_cache: std::sync::Mutex, /// Multi-token-prediction head and driver, present only when this model was given a /// paired MTP GGUF via [`LlamaCppModel::enable_mtp`]. `None` for every other model, /// and every MTP code path is gated on this being `Some` — same shape as `mtmd_ctx`. @@ -485,8 +562,8 @@ impl Drop for LlamaCppModel { // mtmd holds no ownership over the llama model/context (it only borrows a // `llama_model*` at init and a `llama_context*` per eval call), but free it // first regardless, before either of the resources it borrowed goes away. - if let Some(mtmd_ctx) = self.mtmd_ctx { - unsafe { ffi::mtmd_free(mtmd_ctx.as_ptr()) }; + if let Some(pool) = &self.mtmd_pool { + pool.free_all(); } // LoRA adapters borrow the model's weights (loaded via `llama_adapter_lora_init(model, ..)`) // but are otherwise independent objects — free them before the model itself. @@ -521,6 +598,7 @@ impl LlamaCppModel { tensor_split: &[f32], moe_offload_cpu: bool, mmproj_path: Option<&std::path::Path>, + vision_contexts: usize, lora_modules: &[(String, std::path::PathBuf, f32)], reranking: bool, rs_rollback: u32, @@ -867,8 +945,12 @@ impl LlamaCppModel { // Vision/multimodal: load the paired mmproj GGUF via mtmd, if given. A bad // pairing (wrong architecture, corrupt file) fails loudly here rather than - // producing garbage output at inference time. - let mtmd_ctx = match mmproj_path { + // producing garbage output at inference time. `--vision-contexts N` loads N + // contexts into a checkout pool so concurrent requests CLIP-encode in + // parallel; the first context failing is fatal (as before), while a later + // one failing only shrinks the pool — the feature still works, just with + // less parallelism, and a warning says so. + let mtmd_pool = match mmproj_path { Some(p) => { let mmproj_cstr = CString::new( p.to_str() @@ -878,18 +960,42 @@ impl LlamaCppModel { let marker_cstr = CString::new(crate::engine::model::MEDIA_MARKER) .expect("MEDIA_MARKER is a valid C string"); mtmd_params.media_marker = marker_cstr.as_ptr(); - let raw = unsafe { - ffi::mtmd_init_from_file(mmproj_cstr.as_ptr(), model.as_ptr(), mtmd_params) - }; - // marker_cstr/mmproj_cstr only need to outlive the call above — mtmd - // copies both into its own storage during init. - Some(NonNull::new(raw).ok_or_else(|| { - unsafe { ffi::llama_free(ctx.as_ptr()) }; - unsafe { ffi::llama_model_free(model.as_ptr()) }; - anyhow!( - "mtmd_init_from_file failed for {p:?} — check the mmproj matches this model's architecture" - ) - })?) + let n_contexts = vision_contexts.max(1); + let mut pool_ctxs: Vec> = Vec::new(); + for i in 0..n_contexts { + let raw = unsafe { + ffi::mtmd_init_from_file(mmproj_cstr.as_ptr(), model.as_ptr(), mtmd_params) + }; + // marker_cstr/mmproj_cstr only need to outlive the call above — mtmd + // copies both into its own storage during init. + match NonNull::new(raw) { + Some(ptr) => pool_ctxs.push(ptr), + None if i == 0 => { + unsafe { ffi::llama_free(ctx.as_ptr()) }; + unsafe { ffi::llama_model_free(model.as_ptr()) }; + return Err(anyhow!( + "mtmd_init_from_file failed for {p:?} — check the mmproj matches this model's architecture" + )); + } + None => { + tracing::warn!( + context_idx = i, + "failed to create additional mtmd context (likely out of memory); \ + continuing with a pool of {}", + pool_ctxs.len() + ); + break; + } + } + } + let n_embd_inp = unsafe { ffi::llama_model_n_embd_inp(model.as_ptr()) } as usize; + if pool_ctxs.len() > 1 { + tracing::info!( + count = pool_ctxs.len(), + "created mtmd context pool for parallel CLIP encoding" + ); + } + Some(MtmdPool::new(pool_ctxs, n_embd_inp)) } None => None, }; @@ -910,8 +1016,8 @@ impl LlamaCppModel { let adapter: &NonNull = adapter; unsafe { ffi::llama_adapter_lora_free(adapter.as_ptr()) }; } - if let Some(mtmd_ctx) = mtmd_ctx { - unsafe { ffi::mtmd_free(mtmd_ctx.as_ptr()) }; + if let Some(pool) = &mtmd_pool { + pool.free_all(); } unsafe { ffi::llama_free(ctx.as_ptr()) }; unsafe { ffi::llama_model_free(model.as_ptr()) }; @@ -944,7 +1050,10 @@ impl LlamaCppModel { chat_env: std::sync::OnceLock::new(), grammars: dashmap::DashMap::new(), decode_bisection_retries: std::sync::atomic::AtomicU64::new(0), - mtmd_ctx, + mtmd_pool, + clip_cache: std::sync::Mutex::new(lru::LruCache::new( + std::num::NonZeroUsize::new(CLIP_CACHE_ENTRIES).expect("cap is non-zero"), + )), #[cfg(fox_mtp)] mtp: None, lora_adapters, @@ -1299,8 +1408,11 @@ impl LlamaCppModel { grammars: dashmap::DashMap::new(), decode_bisection_retries: std::sync::atomic::AtomicU64::new(0), // bench-kv compares KV cache types, not vision/LoRA; the original - // instance (if any) keeps ownership of its mtmd context and adapters. - mtmd_ctx: None, + // instance (if any) keeps ownership of its mtmd contexts and adapters. + mtmd_pool: None, + clip_cache: std::sync::Mutex::new(lru::LruCache::new( + std::num::NonZeroUsize::new(CLIP_CACHE_ENTRIES).expect("cap is non-zero"), + )), #[cfg(fox_mtp)] mtp: None, rollback_budget: unsafe { @@ -1487,7 +1599,7 @@ impl Model for LlamaCppModel { } fn supports_vision(&self) -> bool { - self.mtmd_ctx.is_some() + self.mtmd_pool.is_some() } fn lora_adapter_names(&self) -> Vec { diff --git a/src/engine/model/llama_cpp/stub.rs b/src/engine/model/llama_cpp/stub.rs index 7f2d2c8..f572258 100644 --- a/src/engine/model/llama_cpp/stub.rs +++ b/src/engine/model/llama_cpp/stub.rs @@ -31,11 +31,12 @@ impl LlamaCppModel { tensor_split: &[f32], moe_offload_cpu: bool, mmproj_path: Option<&std::path::Path>, + vision_contexts: usize, lora_modules: &[(String, std::path::PathBuf, f32)], reranking: bool, rs_rollback: u32, ) -> Result { - let _ = rs_rollback; + let _ = (rs_rollback, vision_contexts); let _ = ( model_path, max_batch_size, diff --git a/src/engine/model/llama_cpp/vocab.rs b/src/engine/model/llama_cpp/vocab.rs index a5c204b..e8849a7 100644 --- a/src/engine/model/llama_cpp/vocab.rs +++ b/src/engine/model/llama_cpp/vocab.rs @@ -153,8 +153,9 @@ impl LlamaCppModel { tools: Option<&serde_json::Value>, images: &[Vec], ) -> Result { - let mtmd_ctx = self - .mtmd_ctx + let pool = self + .mtmd_pool + .as_ref() .ok_or_else(|| anyhow!("model has no vision support (no mmproj loaded)"))?; // Same add_special/parse_special split as build_prompt_tokens_impl: a real @@ -179,57 +180,136 @@ impl LlamaCppModel { parse_special, }; - // mtmd_helper_bitmap_init_from_buf decodes the raw image bytes (bundled - // stb_image). We own the resulting bitmaps until mtmd_tokenize returns — - // it only borrows the pointers to compute chunks, never takes ownership - // (mirrors mtmd-cli.cpp's own usage: bitmaps are freed right after tokenize). - let free_bitmap = |w: &ffi::mtmd_helper_bitmap_wrapper| { - if !w.bitmap.is_null() { - unsafe { ffi::mtmd_bitmap_free(w.bitmap) }; + // Everything below holds a pooled mtmd context checked out for exclusive + // use — decode + tokenize + CLIP encode — and must release it on every + // exit path, hence the closure. With `--vision-contexts N` this is what + // lets N requests preprocess their images concurrently. + let mtmd_nn = pool.acquire(); + let result = (|| -> Result { + let mtmd_ptr = mtmd_nn.as_ptr(); + + // mtmd_helper_bitmap_init_from_buf decodes the raw image bytes (bundled + // stb_image). We own the resulting bitmaps until mtmd_tokenize returns — + // it only borrows the pointers to compute chunks, never takes ownership + // (mirrors mtmd-cli.cpp's own usage: bitmaps are freed right after tokenize). + let free_bitmap = |w: &ffi::mtmd_helper_bitmap_wrapper| { + if !w.bitmap.is_null() { + unsafe { ffi::mtmd_bitmap_free(w.bitmap) }; + } + }; + let mut wrappers = Vec::with_capacity(images.len()); + for img in images { + let wrapper = unsafe { + ffi::mtmd_helper_bitmap_init_from_buf( + mtmd_ptr, + img.as_ptr(), + img.len(), + false, // placeholder + ) + }; + if wrapper.bitmap.is_null() { + wrappers.iter().for_each(free_bitmap); + anyhow::bail!( + "failed to decode image ({} bytes) — unsupported or corrupt format", + img.len() + ); + } + wrappers.push(wrapper); } - }; - let mut wrappers = Vec::with_capacity(images.len()); - for img in images { - let wrapper = unsafe { - ffi::mtmd_helper_bitmap_init_from_buf( - mtmd_ctx.as_ptr(), - img.as_ptr(), - img.len(), - false, // placeholder + let mut bitmap_ptrs: Vec<*const ffi::mtmd_bitmap> = + wrappers.iter().map(|w| w.bitmap as *const _).collect(); + + let chunks_ptr = unsafe { ffi::mtmd_input_chunks_init() }; + let ret = unsafe { + ffi::mtmd_tokenize( + mtmd_ptr, + chunks_ptr, + &input_text, + bitmap_ptrs.as_mut_ptr(), + bitmap_ptrs.len(), ) }; - if wrapper.bitmap.is_null() { - wrappers.iter().for_each(free_bitmap); - anyhow::bail!( - "failed to decode image ({} bytes) — unsupported or corrupt format", - img.len() - ); + wrappers.iter().for_each(free_bitmap); + + if ret != 0 { + unsafe { ffi::mtmd_input_chunks_free(chunks_ptr) }; + anyhow::bail!("mtmd_tokenize failed (code {ret})"); } - wrappers.push(wrapper); - } - let mut bitmap_ptrs: Vec<*const ffi::mtmd_bitmap> = - wrappers.iter().map(|w| w.bitmap as *const _).collect(); - - let chunks_ptr = unsafe { ffi::mtmd_input_chunks_init() }; - let ret = unsafe { - ffi::mtmd_tokenize( - mtmd_ctx.as_ptr(), - chunks_ptr, - &input_text, - bitmap_ptrs.as_mut_ptr(), - bitmap_ptrs.len(), - ) - }; - wrappers.iter().for_each(free_bitmap); - if ret != 0 { - unsafe { ffi::mtmd_input_chunks_free(chunks_ptr) }; - anyhow::bail!("mtmd_tokenize failed (code {ret})"); - } + // From here `chunks` owns chunks_ptr; any error below frees it via Drop. + let mut chunks = unsafe { + crate::engine::model::MultimodalChunks::from_raw( + chunks_ptr as *mut std::ffi::c_void, + ) + }; + + // CLIP-encode every media chunk now, on this pooled context, so prefill + // can decode from the stored embeddings instead of encoding while it + // holds the llama context lock (see `do_prefill_multimodal`). Repeated + // images — chat clients re-send the whole history's images every turn — + // are served from the per-model CLIP cache, keyed by the raw bytes. + let n_chunks = unsafe { ffi::mtmd_input_chunks_size(chunks_ptr) }; + let media_chunk_indices: Vec = (0..n_chunks) + .filter(|&i| { + let chunk = unsafe { ffi::mtmd_input_chunks_get(chunks_ptr, i) }; + let ty = unsafe { ffi::mtmd_input_chunk_get_type(chunk) }; + ty != ffi::mtmd_input_chunk_type_MTMD_INPUT_CHUNK_TYPE_TEXT + }) + .collect(); + // The cache key maps the k-th media chunk back to `images[k]`; that + // holds when mtmd produced exactly one chunk per bitmap. If it didn't + // (an architecture that splits one image into several chunks), encode + // without caching rather than key a chunk by the wrong image's bytes. + let cache_keys: Option> = (media_chunk_indices.len() == images.len()) + .then(|| { + images + .iter() + .map(|b| { + use std::hash::{Hash, Hasher}; + let mut h = ahash::AHasher::default(); + b.hash(&mut h); + (h.finish(), b.len()) + }) + .collect() + }); + + let mut embeddings: Vec<(usize, std::sync::Arc>)> = + Vec::with_capacity(media_chunk_indices.len()); + for (k, &i) in media_chunk_indices.iter().enumerate() { + let key = cache_keys.as_ref().map(|keys| keys[k]); + if let Some(key) = key { + let mut cache = self.clip_cache.lock().expect("clip cache lock poisoned"); + if let Some(hit) = cache.get(&key) { + embeddings.push((i, std::sync::Arc::clone(hit))); + continue; + } + } - Ok(unsafe { - crate::engine::model::MultimodalChunks::from_raw(chunks_ptr as *mut std::ffi::c_void) - }) + let chunk = unsafe { ffi::mtmd_input_chunks_get(chunks_ptr, i) }; + let ret = unsafe { ffi::mtmd_encode_chunk(mtmd_ptr, chunk) }; + if ret != 0 { + anyhow::bail!("mtmd_encode_chunk failed (code {ret})"); + } + let n_tokens = unsafe { ffi::mtmd_input_chunk_get_n_tokens(chunk) }; + let embd_ptr = unsafe { ffi::mtmd_get_output_embd(mtmd_ptr) }; + if embd_ptr.is_null() { + anyhow::bail!("mtmd_get_output_embd returned null after encode"); + } + let embd_len = pool.n_embd_inp * n_tokens; + let embd = std::sync::Arc::new( + unsafe { std::slice::from_raw_parts(embd_ptr, embd_len) }.to_vec(), + ); + if let Some(key) = key { + let mut cache = self.clip_cache.lock().expect("clip cache lock poisoned"); + cache.put(key, std::sync::Arc::clone(&embd)); + } + embeddings.push((i, embd)); + } + chunks.set_image_embeddings(embeddings); + Ok(chunks) + })(); + pool.release(mtmd_nn); + result } pub(super) fn token_to_piece_impl(&self, token: i32) -> Result { diff --git a/src/engine/model/mod.rs b/src/engine/model/mod.rs index a5ec188..c4d45cc 100644 --- a/src/engine/model/mod.rs +++ b/src/engine/model/mod.rs @@ -66,6 +66,15 @@ pub const MEDIA_MARKER: &str = "<__fox_media__>"; #[derive(Debug)] pub struct MultimodalChunks { inner: std::sync::Arc, + /// Pre-encoded CLIP embeddings for the prompt's image/audio chunks, keyed by + /// chunk index: `n_tokens × n_embd_inp` floats per entry, produced at tokenize + /// time on a pooled mtmd context (or served from the per-model CLIP cache). + /// When every non-text chunk has an entry here, prefill decodes from these via + /// `mtmd_helper_decode_image_chunk` instead of re-encoding inside the atomic + /// `mtmd_helper_eval_chunks` call — the encode then never runs while the + /// llama context lock is held. Empty = prefill falls back to the atomic path. + /// `Arc` per entry so a CLIP-cache hit shares the buffer instead of copying it. + image_embeddings: Vec<(usize, std::sync::Arc>)>, } // The field is only read in `#[cfg(not(fox_stub))]` code (Drop, n_positions) — @@ -89,6 +98,7 @@ impl Clone for MultimodalChunks { fn clone(&self) -> Self { Self { inner: self.inner.clone(), + image_embeddings: self.image_embeddings.clone(), } } } @@ -101,6 +111,7 @@ impl MultimodalChunks { pub(crate) unsafe fn from_raw(ptr: *mut std::ffi::c_void) -> Self { Self { inner: std::sync::Arc::new(RawChunksPtr(ptr)), + image_embeddings: Vec::new(), } } @@ -109,6 +120,30 @@ impl MultimodalChunks { self.inner.0 } + /// Attach pre-encoded CLIP embeddings, keyed by chunk index (see the field doc). + #[cfg(not(fox_stub))] + pub(crate) fn set_image_embeddings( + &mut self, + embeddings: Vec<(usize, std::sync::Arc>)>, + ) { + self.image_embeddings = embeddings; + } + + /// Pre-encoded embedding for the chunk at `idx`, if tokenize produced one. + #[cfg(not(fox_stub))] + pub(crate) fn embedding_for_chunk(&self, idx: usize) -> Option<&std::sync::Arc>> { + self.image_embeddings + .iter() + .find(|(i, _)| *i == idx) + .map(|(_, e)| e) + } + + /// Whether this prompt carries pre-encoded embeddings (the split prefill path). + #[cfg(not(fox_stub))] + pub(crate) fn has_image_embeddings(&self) -> bool { + !self.image_embeddings.is_empty() + } + /// Total KV positions this multimodal prompt will occupy (text + image /// tokens combined; for M-RoPE architectures this can differ from the raw /// token count) — used for scheduler block accounting exactly like diff --git a/src/mcp/mod.rs b/src/mcp/mod.rs index 8994113..eb35e6d 100644 --- a/src/mcp/mod.rs +++ b/src/mcp/mod.rs @@ -396,16 +396,8 @@ impl McpServer { let engine = &entry.engine; let messages = vec![("user".to_string(), prompt.to_string())]; - let prompt_text = engine.apply_chat_template(&messages).unwrap_or_else(|_| { - messages - .iter() - .map(|(r, c)| format!("{r}: {c}")) - .collect::>() - .join("\n") - }); - let tokens = engine - .tokenize(&prompt_text) + .build_prompt_tokens(&messages, false, None) .map_err(|e| (INTERNAL_ERROR, format!("tokenize failed: {e}")))?; let text = run_inference(engine, tokens, max_tokens, temperature).await?; @@ -455,16 +447,8 @@ impl McpServer { .map_err(|e| (INTERNAL_ERROR, format!("failed to load model: {e}")))?; let engine = &entry.engine; - let prompt_text = engine.apply_chat_template(&messages).unwrap_or_else(|_| { - messages - .iter() - .map(|(r, c)| format!("{r}: {c}")) - .collect::>() - .join("\n") - }); - let tokens = engine - .tokenize(&prompt_text) + .build_prompt_tokens(&messages, false, None) .map_err(|e| (INTERNAL_ERROR, format!("tokenize failed: {e}")))?; let text = run_inference(engine, tokens, max_tokens, temperature).await?; diff --git a/src/model_registry/config.rs b/src/model_registry/config.rs index ac19737..6fb126b 100644 --- a/src/model_registry/config.rs +++ b/src/model_registry/config.rs @@ -66,6 +66,11 @@ pub struct RegistryConfig { /// whatever model is currently loaded. A mismatched pairing (mmproj for a /// different architecture) fails at load time rather than corrupting output. pub mmproj: Option, + /// Number of mtmd contexts loaded per vision model (`--vision-contexts`). + /// 1 = single context (concurrent encodes serialize on it); N>1 = a checkout + /// pool, so up to N requests CLIP-encode their images in parallel at roughly + /// one mmproj's VRAM footprint per extra context. Meaningless without `mmproj`. + pub vision_contexts: usize, /// Name/path of the multi-token-prediction head GGUF paired with the main model /// (`mtp-*.gguf`), enabling MTP speculative decoding. Like `draft_model` and /// `mmproj`, one global pairing against whichever model is loaded. Only takes @@ -135,6 +140,7 @@ impl RegistryConfig { spec_draft_len: 4, draft_model: None, mmproj: None, + vision_contexts: 1, mtp_model: None, lora_modules: Vec::new(), primary_model: None, diff --git a/src/model_registry/loader.rs b/src/model_registry/loader.rs index 20105b1..a34af09 100644 --- a/src/model_registry/loader.rs +++ b/src/model_registry/loader.rs @@ -53,6 +53,7 @@ async fn load_draft_model( &tensor_split, moe_offload_cpu, None, // mmproj_path — draft models are text-only speculation proposers + 1, // vision_contexts — no mmproj, so no pool to size &[], // lora_modules — adapters apply to the primary model, not the draft false, // reranking — a draft model only ever proposes tokens 0, // rs_rollback — the draft holds no reusable prefix of its own @@ -102,6 +103,7 @@ pub(super) async fn load_model( let split_mode = cfg.split_mode; let tensor_split = cfg.tensor_split.clone(); let moe_offload_cpu = cfg.moe_offload_cpu; + let vision_contexts = cfg.vision_contexts; // Estimate VRAM requirement before attempting to load. // Heuristic: file_size × 1.8 covers weights + overhead. Warn early so the @@ -140,6 +142,7 @@ pub(super) async fn load_model( &tensor_split, moe_offload_cpu, mmproj.as_deref(), + vision_contexts, &lora_modules, reranking, rs_rollback, diff --git a/src/model_registry/mod.rs b/src/model_registry/mod.rs index 9f2bf77..acd6e18 100644 --- a/src/model_registry/mod.rs +++ b/src/model_registry/mod.rs @@ -527,6 +527,7 @@ mod tests { tensor_split: vec![], moe_offload_cpu: false, mmproj: None, + vision_contexts: 1, mtp_model: None, lora_modules: Vec::new(), primary_model: None, From 5f5a665ad40de6cb7ffca7a1fa9648c6b5e47b16 Mon Sep 17 00:00:00 2001 From: CrimsonMartin Date: Fri, 28 Aug 2026 21:46:32 +0000 Subject: [PATCH 5/6] fix: allow result_large_err on load_model_or_respond (clippy 1.98) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CI runs the latest stable clippy, and 1.98's new result_large_err lint flags this function's ~128-byte Err. That Err is the finished HTTP response a handler returns — built at most once per rejected request and consumed immediately — so boxing it would add an unbox at every call site to save a move that is on no hot path. Documented allow instead. --- src/api/error.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/api/error.rs b/src/api/error.rs index 87c7994..448f7e7 100644 --- a/src/api/error.rs +++ b/src/api/error.rs @@ -113,6 +113,12 @@ impl IntoResponse for AppError { /// names a configured adapter (`--lora-modules`) instead of a real model/alias — /// see `ModelRegistry::resolve_for_request`. Callers that don't apply LoRA /// (embeddings) can just discard the second element. +// `result_large_err` (new in clippy 1.98) wants the ~128-byte `Response` boxed, +// but the Err variant here IS the finished HTTP response a handler returns — +// built at most once per rejected request, then consumed immediately. Boxing it +// would add an unbox at every `Err(r) => return r` call site to save a move +// that is not on any hot path. +#[allow(clippy::result_large_err)] pub async fn load_model_or_respond( registry: &ModelRegistry, model: &str, From 578af94717d2b341d7a8010a54bf504b88dc01b6 Mon Sep 17 00:00:00 2001 From: CrimsonMartin Date: Fri, 28 Aug 2026 21:54:25 +0000 Subject: [PATCH 6/6] docs: generalize MCP IDE client examples in module comment --- src/cli/mcp.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cli/mcp.rs b/src/cli/mcp.rs index 8bcca9d..bc6bdc3 100644 --- a/src/cli/mcp.rs +++ b/src/cli/mcp.rs @@ -1,5 +1,5 @@ // `fox mcp` — start an MCP (Model Context Protocol) server over stdio. -// Designed for IDE integration (Cursor, VS Code, Claude Code). +// Designed for IDE integration (Cursor, VS Code, and other MCP clients). use std::path::PathBuf; use std::sync::Arc;