diff --git a/cmake/arm64-windows-msvc-cuda.cmake b/cmake/arm64-windows-msvc-cuda.cmake new file mode 100644 index 000000000..370f2b3d2 --- /dev/null +++ b/cmake/arm64-windows-msvc-cuda.cmake @@ -0,0 +1,26 @@ +# Used to cross-compile ggml-cuda for Windows ARM64 on an x64 Windows host. +set( CMAKE_SYSTEM_NAME Windows ) +set( CMAKE_SYSTEM_PROCESSOR arm64 ) + +if ( DEFINED CUDAToolkit_ROOT ) + file( TO_CMAKE_PATH "${CUDAToolkit_ROOT}" CUDA_ROOT ) +elseif ( DEFINED ENV{CUDA_PATH} ) + file( TO_CMAKE_PATH "$ENV{CUDA_PATH}" CUDA_ROOT ) +else() + message( FATAL_ERROR "Set CUDAToolkit_ROOT or CUDA_PATH to a Windows CUDA Toolkit with ARM64 target libraries" ) +endif() + +if ( DEFINED ENV{VCToolsInstallDir} ) + file( TO_CMAKE_PATH "$ENV{VCToolsInstallDir}" MSVC_TOOLS_ROOT ) + set( CMAKE_CUDA_HOST_COMPILER "${MSVC_TOOLS_ROOT}/bin/Hostx64/arm64/cl.exe" CACHE FILEPATH "" ) +endif() + +set( CMAKE_CUDA_COMPILER "${CUDA_ROOT}/bin/nvcc.exe" CACHE FILEPATH "" ) +set( CMAKE_CUDA_FLAGS_INIT "-target-dir=arm64" ) + +# FindCUDAToolkit selects lib/x64 from the host architecture on Windows. +set( CUDA_CUDART "${CUDA_ROOT}/lib/arm64/cudart.lib" CACHE FILEPATH "" ) +set( CUDA_cudart_LIBRARY "${CUDA_ROOT}/lib/arm64/cudart.lib" CACHE FILEPATH "" ) +set( CUDA_cublas_LIBRARY "${CUDA_ROOT}/lib/arm64/cublas.lib" CACHE FILEPATH "" ) +set( CUDA_cublasLt_LIBRARY "${CUDA_ROOT}/lib/arm64/cublasLt.lib" CACHE FILEPATH "" ) +set( CUDA_cuda_driver_LIBRARY "${CUDA_ROOT}/lib/arm64/cuda.lib" CACHE FILEPATH "" ) diff --git a/common/arg.cpp b/common/arg.cpp index 6a216eb65..a5bfc75c7 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -3313,8 +3313,9 @@ common_params_context common_params_parser_init(common_params & params, llama_ex {"--tools-runtime"}, "OPTION", "experimental: run tools in a separate runtime environment (default: none, use host environment)\n" "available options:\n" - " 'docker:': spin up a new Docker container and reuse it for all invocations, clean up on server exit\n" - " 'docker-container:': use an existing Docker container by ID, won't stop on server exit\n", + " 'docker:', 'podman:': spin up a new container and reuse it for all invocations, clean up on server exit\n" + " 'docker-container:', 'podman-container:': use an existing container by ID, won't stop on server exit\n" + " 'ssh:': run tools on a remote POSIX host over SSH, key-based auth and a trusted host key are required\n", [](common_params & params, const std::string & value) { params.server_tools_runtime = value; } diff --git a/common/chat.cpp b/common/chat.cpp index aabc8e659..0e0b5691b 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -3100,6 +3100,151 @@ static common_chat_params common_chat_params_init_minicpm5(const common_chat_tem return data; } +// An assistant turn is rendered as one or more messages, each +// "<|start|>assistant to=<|message|>{content}{END}" where END is +// <|eom|> (more messages follow) or <|eot|> (end of turn): +// - chain-of-thought: to=self, terminated by <|eom|> +// - final answer: to=user, terminated by <|eot|> +// The generation prompt is just "<|start|>assistant"; the model emits its own +// " to=...<|message|>". +static common_chat_params common_chat_params_init_muse_glimmer(const common_chat_template & tmpl, + const autoparser::generation_params & inputs) { + common_chat_params data; + + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.generation_prompt = "<|start|>assistant"; + data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; + data.supports_thinking = true; + + data.preserved_tokens = { + "<|start|>", "<|message|>", "<|eom|>", "<|eot|>", + // ATEM tool-call markup emitted on " to=" turns. + "", "", + "", "", + }; + + data.message_delimiters = { + { COMMON_CHAT_ROLE_ASSISTANT, "<|start|>assistant" }, + { COMMON_CHAT_ROLE_USER, "<|start|>user" }, + { COMMON_CHAT_ROLE_SYSTEM, "<|start|>system" }, + { COMMON_CHAT_ROLE_TOOL, "<|start|>tool" }, + }; + + if (inputs.has_continuation()) { + const auto & msg = inputs.continue_msg; + + data.generation_prompt = "<|start|>assistant to=self<|message|>" + msg.reasoning_content; + if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { + data.generation_prompt += "<|eom|><|start|>assistant to=user<|message|>" + msg.render_content(); + } + + data.prompt += data.generation_prompt; + } + + auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; + + auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); + // Constrained grammar whenever tools are offered. + auto include_grammar = has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE; + + auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { + auto start = p.rule("start", p.literal("<|start|>assistant")); + + if (!extract_reasoning && !include_grammar) { + return start + p.content(p.rest()); + } + + if (extract_reasoning) { + p.rule("analysis", p.literal(" to=self<|message|>") + p.reasoning(p.until("<|eom|>")) + p.literal("<|eom|>")); + } else { + p.rule("analysis", p.literal(" to=self<|message|>") + p.content(p.until("<|eom|>")) + p.literal("<|eom|>")); + } + auto analysis = p.ref("analysis"); + + auto recipient = p.optional(p.literal(" to=user")); + auto final_msg = p.rule("final", recipient + p.literal("<|message|>") + p.content(p.until("<|eot|>"))); + + if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) { + auto string_value = p.ac( + p.tool_arg_string_value(p.until("")) + p.tool_arg_close(p.literal("")), + ""); + + auto tool_choice = p.choice(); + foreach_function(inputs.tools, [&](const json & tool) { + const auto & function = tool.at("function"); + const std::string name = function.at("name"); + auto params = function.contains("parameters") ? function.at("parameters") : json::object(); + + auto args = p.eps(); + if (params.contains("properties") && params.at("properties").is_object() && !params.at("properties").empty()) { + auto schema_info = common_schema_info(); + schema_info.resolve_refs(params); + + auto arg_choice = p.choice(); + for (const auto & [prop_name, prop_schema] : params.at("properties").items()) { + auto value_parser = p.eps(); + if (schema_info.resolves_to_string(prop_schema)) { + value_parser = string_value; + } else { + value_parser = p.tool_arg_json_value( + p.schema(p.json(), "tool-" + name + "-arg-" + prop_name + "-schema", prop_schema, false)) + + p.tool_arg_close(p.literal("")); + } + + auto arg_rule = p.tool_arg( + p.tool_arg_open(p.literal("")) + + value_parser); + + arg_choice |= arg_rule; + } + args = p.zero_or_more(arg_choice + p.space()); + } + + auto tool_parser = p.tool( + p.tool_open(p.literal(" to=") + p.until("<|message|>") + + p.literal("<|message|>") + p.space() + + p.literal("") + p.space()) + << p.tool_args(args) + << p.tool_close(p.literal("") + p.space() + p.literal(""))); + + tool_choice |= p.rule("tool-" + name, tool_parser); + }); + + auto tool_calls = inputs.parallel_tool_calls + ? p.trigger_rule("tool-call", tool_choice + p.zero_or_more(p.literal("<|eom|>") + start + tool_choice)) + : p.trigger_rule("tool-call", tool_choice); + + + if (inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED) { + return p.zero_or_more(start + analysis) + start + tool_calls; + } + return p.zero_or_more(start + analysis) + start + (tool_calls | final_msg); + } + + return p.zero_or_more(start + analysis) + start + final_msg; + }); + + data.parser = parser.save(); + + if (include_grammar) { + data.grammar_lazy = inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_REQUIRED; + data.grammar = build_grammar([&](const common_grammar_builder & builder) { + foreach_function(inputs.tools, [&](const json & tool) { + const auto & function = tool.at("function"); + auto schema = function.contains("parameters") ? function.at("parameters") : json::object(); + builder.resolve_refs(schema); + }); + parser.build_grammar(builder, data.grammar_lazy); + }); + data.grammar_triggers = { + { COMMON_GRAMMAR_TRIGGER_TYPE_PATTERN, + "<\\|start\\|>assistant( to=(?!self<\\|message\\|>)(?!user<\\|message\\|>)[^<]*?<\\|message\\|>)" }, + }; + } + + return data; +} + static json common_chat_extra_context() { json ctx = json::object(); std::chrono::system_clock::time_point now = std::chrono::system_clock::now(); @@ -3128,6 +3273,12 @@ std::optional common_chat_try_specialized_template( return common_chat_params_init_gpt_oss(tmpl, params); } + // Muse Glimmer format using " to=" recipients and <|eom|>/<|eot|> message terminators. + if (src.find("") != std::string::npos && src.find("<|eom|>") != std::string::npos) { + LOG_DBG("Using specialized template: Muse Glimmer\n"); + return common_chat_params_init_muse_glimmer(tmpl, params); + } + // Functionary v3.2 - uses recipient-based format with >>>recipient\n{content} // Detection: template has ">>>all" for content and ">>>" prefix for tool calls if (src.find(">>>all") != std::string::npos && src.find(">>>${recipient}") != std::string::npos) { diff --git a/common/speculative.cpp b/common/speculative.cpp index 4ce1ea325..02f4c8054 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -1032,7 +1032,14 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { return true; } - if (batch_in.token == nullptr || batch_in.embd != nullptr) { + // Target prefill may contain token IDs or multimodal embeddings. Both + // produce the target-layer features used to seed the draft KV cache, so + // skipping the embedding batches leaves a hole in the draft's cache and + // the next injection fails to initialize. + // TODO: revisit after https://github.com/ggml-org/llama.cpp/pull/24669 is merged + const bool has_tokens = batch_in.token != nullptr; + const bool has_embeddings = batch_in.embd != nullptr; + if (has_tokens == has_embeddings) { return true; } diff --git a/conversion/__init__.py b/conversion/__init__.py index 3b8bebdae..c7d8046c4 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -183,6 +183,8 @@ TEXT_MODEL_MAP: dict[str, str] = { "Olmo3ForCausalLM": "olmo", "OlmoForCausalLM": "olmo", "OlmoeForCausalLM": "olmo", + "MuseGlimmerAssistantModel": "muse_glimmer", + "MuseGlimmerForConditionalGeneration": "muse_glimmer", "OpenELMForCausalLM": "openelm", "OrionForCausalLM": "orion", "PLMForCausalLM": "plm", @@ -298,6 +300,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = { "MiniCPMV4_6ForConditionalGeneration": "minicpm", "Mistral3ForConditionalGeneration": "llava", "NemotronH_Nano_VL_V2": "nemotron", + "MuseGlimmerForConditionalGeneration": "muse_glimmer", "PaddleOCRVisionModel": "ernie", "Phi4ForCausalLMV": "phi", "Qwen2AudioForConditionalGeneration": "ultravox", diff --git a/conversion/muse_glimmer.py b/conversion/muse_glimmer.py new file mode 100644 index 000000000..cc588e832 --- /dev/null +++ b/conversion/muse_glimmer.py @@ -0,0 +1,179 @@ +from __future__ import annotations + +import json +from typing import Any, Iterable, TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from torch import Tensor + +from .base import MmprojModel, ModelBase, TextModel, gguf + + +def _unpermute_for_rope(tensor: "Tensor", n_heads: int) -> "Tensor": + """Invert transformers' `_permute_for_rope`: HF stores Q/K in rotate_half layout, + llama.cpp consumes the interleaved (NORM) layout.""" + if tensor.ndim == 2: + dim1, dim2 = tensor.shape + return tensor.view(n_heads, 2, dim1 // n_heads // 2, dim2).transpose(1, 2).reshape(dim1, dim2) + if tensor.ndim == 1: + (dim1,) = tensor.shape + return tensor.view(n_heads, 2, dim1 // n_heads // 2).transpose(1, 2).reshape(dim1) + raise ValueError(f"_unpermute_for_rope: unexpected shape {tuple(tensor.shape)}") + + +@ModelBase.register("MuseGlimmerForConditionalGeneration") +class MuseGlimmerModel(TextModel): + model_arch = gguf.MODEL_ARCH.MUSE_GLIMMER + + def norm_shift(self, name: str) -> float: + # All four layer norms use 1, the final norm uses 0. + return 1.0 if name.endswith("layernorm.weight") else 0.0 + + def set_vocab(self): + self._set_vocab_gpt2() + + from transformers import AutoTokenizer + tok = AutoTokenizer.from_pretrained(self.dir_model) + eot_id = tok.convert_tokens_to_ids("<|eot|>") + if isinstance(eot_id, int) and eot_id >= 0: + self.gguf_writer.add_eot_token_id(eot_id) + + def set_gguf_parameters(self): + super().set_gguf_parameters() + hparams = self.hparams + + self.gguf_writer.add_final_logit_softcapping(hparams["final_logit_softcapping"]) + self.gguf_writer.add_logit_scale(hparams["output_multiplier"]) + self.gguf_writer.add_sliding_window(hparams["sliding_window"]) + self.gguf_writer.add_sliding_window_pattern([t == "sliding_attention" for t in hparams["layer_types"]]) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + shift = self.norm_shift(name) + if shift != 0.0: + data_torch = data_torch + shift + + # Invert transformers' `_permute_for_rope` on Q/K, we keep ggml's NORM (interleaved) rope + if ".self_attn.q_proj." in name: + data_torch = _unpermute_for_rope(data_torch, int(self.hparams["num_attention_heads"])) + elif ".self_attn.k_proj." in name: + data_torch = _unpermute_for_rope(data_torch, int(self.hparams["num_key_value_heads"])) + + # Synthesize QK-norm weights to absorb qk_scale_factor. + # MuseGlimmer implementation: scaleless RMSNorm followed by qk_scale_factor.. + if bid is not None and name.endswith(f"model.layers.{bid}.self_attn.q_proj.weight"): + head_dim = self.hparams["head_dim"] + q_scale = float(self.hparams["qk_scale_factor"]) + yield ( + self.map_tensor_name(f"model.layers.{bid}.self_attn.q_norm.weight"), + torch.full((head_dim,), q_scale, dtype=torch.float32), + ) + yield ( + self.map_tensor_name(f"model.layers.{bid}.self_attn.k_norm.weight"), + torch.ones((head_dim,), dtype=torch.float32), + ) + + yield from super().modify_tensors(data_torch, name, bid) + + +@ModelBase.register("MuseGlimmerForConditionalGeneration") +class MuseGlimmerVisionModel(MmprojModel): + def get_vision_config(self) -> dict[str, Any] | None: + c = self.global_config.get("vision_config") + if not c: + return None + # MuseGlimmer actually uses dynamic size, initialize with nominal size + image_size = c["pos_emb_height"] * c["patch_size"] * c["merge_size"] + return {**c, "image_size": image_size} + + def set_gguf_parameters(self): + super().set_gguf_parameters() + assert self.hparams_vision is not None + c = self.hparams_vision # enriched vision_config from get_vision_config() + + self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.MUSE_GLIMMER) + self.gguf_writer.add_vision_attention_layernorm_eps(float(c["layer_norm_eps"])) + self.gguf_writer.add_vision_spatial_merge_size(int(c["merge_size"])) + + @classmethod + def filter_tensors(cls, item): + name, gen = item + keep = ("model.vision_tower.", "model.vision_adapter.", "model.vision_projection.") + if not any(name.startswith(k) for k in keep): + return None + return super().filter_tensors((name, gen)) + + # 3-layer projector MLP + _MM_MLP_MAP = { + "model.vision_adapter.fc1": (gguf.MODEL_TENSOR.V_MMPROJ, 0), + "model.vision_adapter.fc2": (gguf.MODEL_TENSOR.V_MMPROJ, 1), + "model.vision_projection": (gguf.MODEL_TENSOR.V_MMPROJ, 2), + } + + def modify_tensors(self, data_torch, name, bid): + assert self.hparams_vision is not None + if ".attn.q_proj." in name or ".attn.k_proj." in name: + n_heads = int(self.hparams_vision["num_attention_heads"]) + data_torch = _unpermute_for_rope(data_torch, n_heads) + # Lay out the pt=2 temporal slabs of the patch embedding as a conv2d for build_inp() + if name.endswith("patch_embedder.patch_embedding.weight"): + n_embd = data_torch.shape[0] + pt = int(self.hparams_vision["patch_temporal"]) + ps = int(self.hparams_vision["patch_size"]) + data_torch = data_torch.view(n_embd, pt, 3, ps, ps).sum(dim=1) # (n_embd, 3, ps, ps) + stem, _, suffix = name.rpartition(".") + if stem in self._MM_MLP_MAP: + tensor_key, idx = self._MM_MLP_MAP[stem] + yield (self.format_tensor_name(tensor_key, bid=idx, suffix="." + suffix), data_torch) + return + yield (self.map_tensor_name(name), data_torch) + + +@ModelBase.register("MuseGlimmerAssistantModel") +class MuseGlimmerAssistantModel(TextModel): + model_arch = gguf.MODEL_ARCH.DFLASH + + def set_vocab(self): + if self.target_model_dir is None: + raise ValueError( + "MuseGlimmerAssistant (DFlash drafter) requires --target-model-dir pointing to the " + "target MuseGlimmer HF directory" + ) + + original_dir = self.dir_model + self.dir_model = self.target_model_dir + + from . import get_model_class + with open(self.target_model_dir / "config.json", "r", encoding="utf-8") as f: + target_arch = json.load(f)["architectures"][0] + target_cls = get_model_class(target_arch) + if target_cls is not type(self): + target_cls.set_vocab(self) # ty: ignore[unresolved-attribute] + else: + super().set_vocab() + + self.dir_model = original_dir + + mask_token_id = self.hparams.get("mask_token_id") + if mask_token_id is not None: + self.gguf_writer.add_mask_token_id(int(mask_token_id)) + + def set_gguf_parameters(self): + super().set_gguf_parameters() + h = self.hparams + + self.gguf_writer.add_block_size(int(h["block_size"])) + + # dflash.target_layers[k] refers to the inputs going into the ith layer, which come from the (i-1)th layer's output. + # The transformers configuration refers to the outputs being recorded. + self.gguf_writer.add_target_layers([int(x) + 1 for x in h["target_layer_ids"]]) + + if h.get("sliding_window") and h.get("layer_types"): + self.gguf_writer.add_sliding_window(int(h["sliding_window"])) + self.gguf_writer.add_sliding_window_pattern([t == "sliding_attention" for t in h["layer_types"]]) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + # DFlash defaults to NEOX (rotate_half) rope, matching transformers HF layout for Q/K, QK-norms + # no permutation needed. + yield (self.map_tensor_name(name), data_torch) diff --git a/conversion/nemotron.py b/conversion/nemotron.py index 0572b42ca..e5075020c 100644 --- a/conversion/nemotron.py +++ b/conversion/nemotron.py @@ -197,6 +197,7 @@ class NemotronHModel(GraniteHybridModel): """Hybrid mamba2/attention model from NVIDIA""" model_arch = gguf.MODEL_ARCH.NEMOTRON_H is_moe: bool = False + supports_mtp_export = True def __init__(self, *args, **kwargs): # We have to determine the correct model architecture (MoE vs non-MoE) before @@ -236,6 +237,25 @@ class NemotronHModel(GraniteHybridModel): self._ssm_layers = [i for i, val in enumerate(pattern) if val == "mamba"] self._mlp_layers = [i for i, val in enumerate(pattern) if val == "moe"] + # `--no-mtp` drops it entirely; `--mtp` exports only the MTP head + self._mtp_bid: int | None = None + if self.is_moe and not self.no_mtp: + n_nextn = self.hparams.get("num_nextn_predict_layers", 0) or 0 + if n_nextn > 0: + assert n_nextn == 1, ( + "NemotronH MTP conversion currently supports num_nextn_predict_layers == 1" + ) + self._mtp_bid = self.block_count + self.block_count += 1 + # The folded MTP block carries both an attention sub-layer and a + # MoE sub-layer, so register it as both so the per-layer metadata arrays cover it + self._attn_layers.append(self._mtp_bid) + self._mlp_layers.append(self._mtp_bid) + self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) + + if self.mtp_only and self._mtp_bid is None: + raise ValueError("--mtp was requested, but this model does not contain a supported MTP head") + def get_attn_layers(self): pattern = self.hparams.get("hybrid_override_pattern") or self.hparams.get("layers_block_type") if pattern is None: @@ -246,6 +266,36 @@ class NemotronHModel(GraniteHybridModel): return [i for i, val in enumerate(pattern) if val == "attention"] + @classmethod + def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: + name, gen = item + if name.startswith("mtp."): + # --no-mtp: drop the MTP head entirely + if cls.no_mtp: + return None + elif cls.mtp_only: + # --mtp: export the MTP head plus the tensors it shares with the target model + keep = name in ( + "backbone.embeddings.weight", + "backbone.norm_f.weight", + "lm_head.weight", + ) + if not keep: + return None + return super().filter_tensors((name, gen)) + + def prepare_metadata(self, vocab_only: bool): + from_dir = self.fname_out.is_dir() + super().prepare_metadata(vocab_only=vocab_only) + + if not self.mtp_only or not from_dir: + return + output_type: str = self.ftype.name.partition("_")[2] + fname_default: str = gguf.naming_convention( + self.metadata.name, self.metadata.basename, self.metadata.finetune, + self.metadata.version, size_label=None, output_type=output_type, model_type=None) + self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf" + def set_gguf_parameters(self): super().set_gguf_parameters() @@ -284,6 +334,10 @@ class NemotronHModel(GraniteHybridModel): if (latent_size := self.hparams.get("moe_latent_size")) is not None: self.gguf_writer.add_moe_latent_size(latent_size) + # MTP head: number of trailing NextN blocks + if self._mtp_bid is not None: + self.gguf_writer.add_nextn_predict_layers(self.hparams["num_nextn_predict_layers"]) + def set_vocab(self): # The NemotronH config uses pattern characters (e.g. '-') that may not # be supported by the installed transformers version. AutoTokenizer @@ -350,15 +404,24 @@ class NemotronHModel(GraniteHybridModel): if not self.is_moe: self.gguf_writer.add_add_bos_token(True) - def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: - if self.is_moe and bid is not None: - # Skip Multi-Token Prediction (MTP) tensors. These are used for - # for speculative decoding but we don't include them in this model - # conversion. See https://github.com/ggml-org/llama.cpp/pull/18886 - if name.startswith("mtp."): - logger.info(f"gguf: Skipping MTP (Speculative) layer: {name}") - return + _MTP_SPECIAL_RENAMES = { + "mtp.layers.0.enorm.weight": "model.layers.{bid}.enorm.weight", + "mtp.layers.0.hnorm.weight": "model.layers.{bid}.hnorm.weight", + "mtp.layers.0.eh_proj.weight": "model.layers.{bid}.eh_proj.weight", + "mtp.layers.1.norm.weight": "model.layers.{bid}.post_attention_layernorm.weight", + "mtp.layers.1.final_layernorm.weight": "model.layers.{bid}.shared_head.norm.weight", + } + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + # mtp.layers.0: NextN input fusion + attention + # mtp.layers.1: MoE + final head norm + if self._mtp_bid is not None and name.startswith(("mtp.layers.0.", "mtp.layers.1.")): + suffix = name.split(".", 3)[3] + bid = self._mtp_bid + renamed = self._MTP_SPECIAL_RENAMES.get(name) + name = renamed.format(bid=bid) if renamed else f"backbone.layers.{bid}.{suffix}" + + if self.is_moe and bid is not None: if name.endswith("mixer.gate.e_score_correction.bias"): yield from ModelBase.modify_tensors(self, data_torch, name, bid) return diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 29957043d..909a8f28e 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -3443,7 +3443,7 @@ static bool ggml_thread_apply_priority(int32_t prio) { return true; } -#elif defined(__gnu_linux__) +#elif defined(__linux__) // TODO: this may not work on BSD, to be verified static bool ggml_thread_apply_affinity(const bool * mask) { diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 6a8e77fbd..0a3cb86b8 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5198,7 +5198,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g return max_bias == 0.0f; } case GGML_OP_ROLL: - if(op->src[0]->type == GGML_TYPE_F32) { + if(op->src[0]->type == GGML_TYPE_F32 && ggml_is_contiguous(op->src[0])) { return true; } return false; diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index 1e0316a9f..fc7d2e727 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1274,8 +1274,9 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te case GGML_OP_ARGSORT: case GGML_OP_TOP_K: case GGML_OP_ARANGE: - case GGML_OP_ROLL: return true; + case GGML_OP_ROLL: + return ggml_is_contiguous(op->src[0]); case GGML_OP_FLASH_ATTN_EXT: // for new head sizes, add checks here if (op->src[0]->ne[0] != 32 && diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 304100f7f..8d0cad59b 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -509,6 +509,7 @@ class MODEL_ARCH(IntEnum): OLMO = auto() OLMO2 = auto() OLMOE = auto() + MUSE_GLIMMER = auto() OPENELM = auto() ARCTIC = auto() DEEPSEEK = auto() @@ -1181,6 +1182,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = { MODEL_ARCH.OLMO: "olmo", MODEL_ARCH.OLMO2: "olmo2", MODEL_ARCH.OLMOE: "olmoe", + MODEL_ARCH.MUSE_GLIMMER: "muse-glimmer", MODEL_ARCH.OPENELM: "openelm", MODEL_ARCH.ARCTIC: "arctic", MODEL_ARCH.DEEPSEEK: "deepseek", @@ -1562,8 +1564,8 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = { MODEL_TENSOR.V_MM_UP: "mm.up", MODEL_TENSOR.V_MM_DOWN: "mm.down", MODEL_TENSOR.V_MM_GATE: "mm.gate", - MODEL_TENSOR.V_MM_MERGER_FC1: "mm.merger.fc1", - MODEL_TENSOR.V_MM_MERGER_FC2: "mm.merger.fc2", + MODEL_TENSOR.V_MM_MERGER_FC1: "mm.merger.fc1", + MODEL_TENSOR.V_MM_MERGER_FC2: "mm.merger.fc2", MODEL_TENSOR.V_TOK_BOI: "v.boi", MODEL_TENSOR.V_TOK_EOI: "v.eoi", MODEL_TENSOR.V_MM_PRE_NORM: "mm.pre_norm", @@ -3331,6 +3333,25 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.FFN_UP_EXP, MODEL_TENSOR.FFN_DOWN_EXP, ], + MODEL_ARCH.MUSE_GLIMMER: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_GATE, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.FFN_PRE_NORM, + MODEL_TENSOR.FFN_POST_NORM, + ], MODEL_ARCH.OPENELM: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, @@ -3846,6 +3867,12 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.FFN_DOWN_SHEXP, MODEL_TENSOR.FFN_UP_SHEXP, MODEL_TENSOR.FFN_EXP_PROBS_B, + # NextN/MTP (draft head) + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, ], MODEL_ARCH.EXAONE: [ MODEL_TENSOR.TOKEN_EMBD, @@ -5160,6 +5187,7 @@ class VisionProjectorType: MIMOVL = "mimovl" MIMO_AUDIO = "mimo_audio" GRANITE4_VISION = "granite4_vision" + MUSE_GLIMMER = "muse-glimmer" # Items here are (block size, type size) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 7892342e4..79d270ab8 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -382,7 +382,7 @@ class TensorNameMap: ), MODEL_TENSOR.ATTN_GATE: ( - "model.layers.{bid}.self_attn.gate_proj", # afmoe + "model.layers.{bid}.self_attn.gate_proj", # afmoe muse-glimmer "model.layers.{bid}.linear_attn.in_proj_z", # qwen3.5 "model.layers.{bid}.self_attn.g_proj", # step3.5 head-wise attention gate ), @@ -1298,10 +1298,12 @@ class TensorNameMap: "encoder.final_layer_norm", # t5 "layer_norm", # neobert "model.hidden_norm", # dflash + "encoder.output_norm_enc", # dflash (transformers MuseGlimmerAssistant) ), MODEL_TENSOR.FC: ( - "model.fc", # dflash + "model.fc", # dflash + "encoder.fc", # dflash (transformers MuseGlimmerAssistant) ), MODEL_TENSOR.DSPARK_MARKOV_W1: ( @@ -1467,6 +1469,7 @@ class TensorNameMap: "vision_tower.patch_embed.patchifier.proj", # dots.ocr "vision_model.conv1", # Step3-VL "model.vision_embedder.patch_dense", # gemma4 unified + "model.vision_tower.patch_embedder.patch_embedding", # muse-glimmer ), MODEL_TENSOR.V_ENC_EMBD_NORM: ( @@ -1534,7 +1537,8 @@ class TensorNameMap: "siglip2.vision_model.encoder.layers.{bid}.self_attn.q_proj", # youtuvl "model.vision_model.transformer.layers.{bid}.self_attn.q_proj", # Deepseek-OCR CLIP, generated "vision_model.model.layers.{bid}.self_attn.q_proj.linear", # gemma4 - "model.qwen2_model.model.model.layers.{bid}.self_attn.q_proj" # Deepseek-OCR-2 qwen2 + "model.qwen2_model.model.model.layers.{bid}.self_attn.q_proj", # Deepseek-OCR-2 qwen2 + "model.vision_tower.layers.{bid}.attn.q_proj", # muse-glimmer ), MODEL_TENSOR.V_ENC_ATTN_Q_NORM: ( @@ -1560,7 +1564,8 @@ class TensorNameMap: "model.vision_model.transformer.layers.{bid}.self_attn.k_proj", # Deepseek-OCR CLIP, generated "siglip2.vision_model.encoder.layers.{bid}.self_attn.k_proj", "vision_model.model.layers.{bid}.self_attn.k_proj.linear", # gemma4 - "model.qwen2_model.model.model.layers.{bid}.self_attn.k_proj" # Deepseek-OCR-2 qwen2 + "model.qwen2_model.model.model.layers.{bid}.self_attn.k_proj", # Deepseek-OCR-2 qwen2 + "model.vision_tower.layers.{bid}.attn.k_proj", # muse-glimmer ), MODEL_TENSOR.V_ENC_ATTN_K_NORM: ( @@ -1586,7 +1591,8 @@ class TensorNameMap: "siglip2.vision_model.encoder.layers.{bid}.self_attn.v_proj", "model.vision_model.transformer.layers.{bid}.self_attn.v_proj", # Deepseek-OCR CLIP, generated "vision_model.model.layers.{bid}.self_attn.v_proj.linear", # gemma4 - "model.qwen2_model.model.model.layers.{bid}.self_attn.v_proj" # Deepseek-OCR-2 qwen2 + "model.qwen2_model.model.model.layers.{bid}.self_attn.v_proj", # Deepseek-OCR-2 qwen2 + "model.vision_tower.layers.{bid}.attn.v_proj", # muse-glimmer ), MODEL_TENSOR.V_ENC_INPUT_NORM: ( @@ -1610,6 +1616,7 @@ class TensorNameMap: "vision_tower.blocks.{bid}.norm1", # dots.ocr "vision_model.transformer.resblocks.{bid}.ln_1", # Step3-VL "model.qwen2_model.model.model.layers.{bid}.input_layernorm", # Deepseek-OCR-2 qwen2 + "model.vision_tower.layers.{bid}.norm1", # muse-glimmer ), MODEL_TENSOR.V_ENC_ATTN_O: ( @@ -1635,6 +1642,7 @@ class TensorNameMap: "vision_model.model.layers.{bid}.self_attn.o_proj.linear", # gemma4 "vision_tower.blocks.{bid}.attn.proj", # dots.ocr "vision_model.transformer.resblocks.{bid}.attn.out_proj", # Step3-VL + "model.vision_tower.layers.{bid}.attn.proj", # muse-glimmer ), MODEL_TENSOR.V_ENC_ATTN_SINKS: ( @@ -1663,6 +1671,7 @@ class TensorNameMap: "vision_tower.blocks.{bid}.norm2", # dots.ocr "vision_model.transformer.resblocks.{bid}.ln_2", # Step3-VL "model.qwen2_model.model.model.layers.{bid}.post_attention_layernorm", # Deepseek-OCR-2 qwen2 + "model.vision_tower.layers.{bid}.norm2", # muse-glimmer ), MODEL_TENSOR.V_ENC_FFN_UP: ( @@ -1687,6 +1696,7 @@ class TensorNameMap: "vision_model.model.layers.{bid}.mlp.up_proj", # gemma4 "vision_model.transformer.resblocks.{bid}.mlp.c_fc", # Step3-VL "model.qwen2_model.model.model.layers.{bid}.mlp.up_proj", # Deepseek-OCR-2 qwen2 + "model.vision_tower.layers.{bid}.mlp.fc1", # muse-glimmer ), MODEL_TENSOR.V_ENC_FFN_GATE: ( @@ -1719,6 +1729,7 @@ class TensorNameMap: "model.qwen2_model.model.model.layers.{bid}.mlp.down_proj" , # Deepseek-OCR-2 qwen2 "vision_model.model.layers.{bid}.mlp.down_proj", # gemma4 "vision_model.transformer.resblocks.{bid}.mlp.c_proj", # Step3-VL + "model.vision_tower.layers.{bid}.mlp.fc2", # muse-glimmer ), MODEL_TENSOR.V_ENC_ATTN_POST_NORM: ( @@ -1753,6 +1764,7 @@ class TensorNameMap: "model.vision_model.pre_layrnorm", # Deepseek-OCR CLIP "vision_tower.patch_embed.patchifier.norm", # dots.ocr "vision_model.ln_pre", # Step3-VL + "model.vision_tower.ln_pre", # muse-glimmer ), MODEL_TENSOR.V_POST_NORM: ( @@ -1766,6 +1778,7 @@ class TensorNameMap: "visual.post_layernorm", # glm4v "siglip2.vision_model.post_layernorm", "model.qwen2_model.model.model.norm", # Deepseek-OCR-2 qwen2 + "model.vision_tower.ln_post", # muse-glimmer ), MODEL_TENSOR.V_MM_POST_NORM: ( diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index b8231ba6f..73fb8b981 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -71,6 +71,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_OLMO, "olmo" }, { LLM_ARCH_OLMO2, "olmo2" }, { LLM_ARCH_OLMOE, "olmoe" }, + { LLM_ARCH_MUSE_GLIMMER, "muse-glimmer" }, { LLM_ARCH_OPENELM, "openelm" }, { LLM_ARCH_ARCTIC, "arctic" }, { LLM_ARCH_DEEPSEEK, "deepseek" }, diff --git a/src/llama-arch.h b/src/llama-arch.h index 47adc3d68..51dfd288c 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -76,6 +76,7 @@ enum llm_arch { LLM_ARCH_OLMO, LLM_ARCH_OLMO2, LLM_ARCH_OLMOE, + LLM_ARCH_MUSE_GLIMMER, LLM_ARCH_OPENELM, LLM_ARCH_ARCTIC, LLM_ARCH_DEEPSEEK, diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index 248d1ed38..abca773a9 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -27,6 +27,7 @@ bool llama_model_saver_supports_arch(llm_arch arch) { case LLM_ARCH_APERTUS: case LLM_ARCH_MIMO2: case LLM_ARCH_STEP35: + case LLM_ARCH_MUSE_GLIMMER: case LLM_ARCH_MELLUM: case LLM_ARCH_LAGUNA: return false; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 1ba795763..78c31ad19 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -128,6 +128,7 @@ #include "models/mistral4.cpp" #include "models/modern-bert.cpp" #include "models/mpt.cpp" +#include "models/muse-glimmer.cpp" #include "models/nanbeige.cpp" #include "models/nemotron-h-moe.cpp" #include "models/nemotron-h.cpp" @@ -180,9 +181,12 @@ #include "models/talkie.cpp" #include "models/wavtokenizer-dec.cpp" #include "models/xverse.cpp" +#include "models/clip.cpp" static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params & params) { switch (arch) { + case LLM_ARCH_CLIP: + return new llama_model_clip(params); case LLM_ARCH_LLAMA: return new llama_model_llama(params); case LLM_ARCH_LLAMA4: @@ -317,6 +321,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_olmo2(params); case LLM_ARCH_OLMOE: return new llama_model_olmoe(params); + case LLM_ARCH_MUSE_GLIMMER: + return new llama_model_muse_glimmer(params); case LLM_ARCH_OPENELM: return new llama_model_openelm(params); case LLM_ARCH_GPTNEOX: @@ -2374,6 +2380,9 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE); + const bool mtp_on_hybrid_nemotron = + params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && arch == LLM_ARCH_NEMOTRON_H_MOE; + if (llm_arch_is_recurrent(arch)) { res = new llama_memory_recurrent( *this, @@ -2384,7 +2393,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, cparams.n_seq_max, cparams.n_rs_seq, nullptr); - } else if (llm_arch_is_hybrid(arch) && !mtp_on_hybrid_qwen) { + } else if (llm_arch_is_hybrid(arch) && !mtp_on_hybrid_qwen && !mtp_on_hybrid_nemotron) { // The main difference between hybrid architectures is the // layer filters, so pick the right one here llama_memory_hybrid::layer_filter_cb filter_attn = nullptr; @@ -2465,7 +2474,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, }; } - if (mtp_on_hybrid_qwen) { + if (mtp_on_hybrid_qwen || mtp_on_hybrid_nemotron) { filter = [&](uint32_t il) { return il >= hparams.n_layer(); }; } @@ -2737,6 +2746,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_DEEPSEEK2OCR: case LLM_ARCH_DEEPSEEK32: case LLM_ARCH_DEEPSEEK4: + case LLM_ARCH_MUSE_GLIMMER: case LLM_ARCH_PLM: case LLM_ARCH_CHATGLM: case LLM_ARCH_GRANITE: diff --git a/src/models/clip.cpp b/src/models/clip.cpp new file mode 100644 index 000000000..537766aeb --- /dev/null +++ b/src/models/clip.cpp @@ -0,0 +1,18 @@ +#include "models.h" + +// Stub to allow llama-quantize to open mmproj GGUFs + +[[noreturn]] +void llama_model_clip::load_arch_hparams(llama_model_loader &) { + GGML_ABORT("CLIP is a quant-only stub; load_arch_hparams should not be called"); +} + +[[noreturn]] +void llama_model_clip::load_arch_tensors(llama_model_loader &) { + GGML_ABORT("CLIP is a quant-only stub; load_arch_tensors should not be called"); +} + +[[noreturn]] +std::unique_ptr llama_model_clip::build_arch_graph(const llm_graph_params &) const { + GGML_ABORT("CLIP has no inference graph via llama_model dispatch; runtime lives in tools/mtmd/clip.cpp"); +} diff --git a/src/models/models.h b/src/models/models.h index a8908da42..923034520 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -386,6 +386,22 @@ struct llama_model_bloom : public llama_model_base { }; +// Quant-only stub for mmproj GGUFs +// none of these are ever called, they only exist to satisfy the llama_model_base interface +struct llama_model_clip : public llama_model_base { + llama_model_clip(const struct llama_model_params & params) : llama_model_base(params) {} + + [[noreturn]] + void load_arch_hparams(llama_model_loader & ml) override; + + [[noreturn]] + void load_arch_tensors(llama_model_loader & ml) override; + + [[noreturn]] + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + + struct llama_model_mpt : public llama_model_base { llama_model_mpt(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; @@ -1028,6 +1044,19 @@ struct llama_model_olmoe : public llama_model_base { }; +struct llama_model_muse_glimmer : public llama_model_base { + llama_model_muse_glimmer(const struct llama_model_params & params) : llama_model_base(params) {} + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + + struct graph : public llm_graph_context { + graph(const llama_model & model, const llm_graph_params & params); + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + + struct llama_model_openelm : public llama_model_base { llama_model_openelm(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; @@ -1461,6 +1490,10 @@ struct llama_model_nemotron_h_moe : public llama_model_nemotron_h { using graph = llama_model_nemotron_h::graph; + struct graph_mtp : public llm_graph_context { + graph_mtp(const llama_model & model, const llm_graph_params & params); + }; + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; }; diff --git a/src/models/muse-glimmer.cpp b/src/models/muse-glimmer.cpp new file mode 100644 index 000000000..0e9415308 --- /dev/null +++ b/src/models/muse-glimmer.cpp @@ -0,0 +1,208 @@ +#include "models.h" + +void llama_model_muse_glimmer::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa); + ml.get_key(LLM_KV_FINAL_LOGIT_SOFTCAPPING, hparams.f_final_logit_softcapping, false); + ml.get_key(LLM_KV_LOGIT_SCALE, hparams.f_logit_scale); + + hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train; + ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false); + + hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; + uint32_t swa_period = 4; + if (ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, swa_period, false)) { + hparams.set_swa_pattern(swa_period); + } else { + ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl, hparams.n_layer()); + } + + switch (hparams.n_layer()) { + case 52: type = LLM_TYPE_30B; break; + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_muse_glimmer::load_arch_tensors(llama_model_loader &) { + LLAMA_LOAD_LOCALS; + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0); + + for (int i = 0; i < n_layer; ++i) { + auto & layer = layers[i]; + + // Pre/post-attention norms (Muse Glimmer's `weight + 1` applied at conversion time). + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), {n_embd}, 0); + + // Q/K/V/O projections. + create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, 0); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, 0); + + // QK-norm. Weights are synthesized at conversion time to absorb `qk_scale_factor`. + layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0); + layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, 0); + + // Attention output gate: sigmoid(gate) * attn_out before o_proj (same as afmoe). + layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_embd_head_k * n_head}, 0); + + // Pre/post-FFN norms (FFN_PRE_NORM is aliased to LLM_TENSOR_FFN_NORM). + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + layer.ffn_post_norm = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), {n_embd}, 0); + + // Dense FFN (unlike afmoe, no MoE branches). + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0); + } +} + +llama_model_muse_glimmer::graph::graph(const llama_model & model, const llm_graph_params & params) + : llm_graph_context(params) { + const int64_t n_embd_head = hparams.n_embd_head_v(); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); + + // Different to f_norm_rms_eps for post-attn / post-FFN norms + const float post_norm_eps = 1e-8f; + + ggml_tensor * cur; + ggml_tensor * inpL; + + inpL = build_inp_embd(model.tok_embd); + inpL = build_norm(inpL, nullptr, nullptr, LLM_NORM_RMS, -1); + cb(inpL, "embd_norm", -1); + + ggml_tensor * inp_pos = build_inp_pos(); + auto * inp_attn = build_attn_inp_kv_iswa(); + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + const float kq_scale = 1.0f / sqrtf(float(n_embd_head)); + + for (int il = 0; il < n_layer; ++il) { + // expose per-layer residual for speculative drafts (see LLM_KV_TARGET_LAYERS). + res->t_layer_inp[il] = inpL; + + const float freq_base_l = model.get_rope_freq_base (cparams, il); + const float freq_scale_l = model.get_rope_freq_scale(cparams, il); + + ggml_tensor * inpSA = inpL; + + // RoPE runs on the SWA layers, NoPE on full ones. + const bool use_rope = hparams.is_swa(il); + + // pre-attention norm (weight+1 folded at conversion time) + cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + + // self-attention: attention output gate around SDPA (afmoe.cpp:147-191) + { + ggml_tensor * attn_inp = cur; // save input for gate computation + + auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, + n_embd_head, n_head, n_head_kv, il); + + // gate = wqkv_gate @ attn_inp (from pre-attn hidden state) + ggml_tensor * gate = build_lora_mm(model.layers[il].wqkv_gate, attn_inp); + cb(gate, "attn_gate_proj", il); + + // QK-norm. attn_q_norm weight was synthesized at conversion to broadcast + // qk_scale_factor across head_dim; attn_k_norm is identity (ones). + Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL, LLM_NORM_RMS, il); + Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, NULL, LLM_NORM_RMS, il); + cb(Qcur, "Qcur_normed", il); + cb(Kcur, "Kcur_normed", il); + + if (use_rope) { + Qcur = ggml_rope_ext( + ctx0, Qcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l, + ext_factor, attn_factor, beta_fast, beta_slow); + cb(Qcur, "Qcur_rope", il); + + Kcur = ggml_rope_ext( + ctx0, Kcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l, + ext_factor, attn_factor, beta_fast, beta_slow); + cb(Kcur, "Kcur_rope", il); + } + + // SDPA. wo is deferred; the gate goes between attn_out and o_proj. + cur = build_attn(inp_attn, + NULL, NULL, NULL, + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); + cb(cur, "attn_out", il); + + gate = ggml_sigmoid(ctx0, gate); + cb(gate, "attn_gate_sig", il); + cur = ggml_mul(ctx0, cur, gate); + cb(cur, "attn_gated", il); + + cur = build_lora_mm(model.layers[il].wo, cur, model.layers[il].wo_s); + cb(cur, "attn_o_proj", il); + } + + cur = ggml_rms_norm(ctx0, cur, post_norm_eps); + cur = ggml_mul(ctx0, cur, model.layers[il].attn_post_norm); + cb(cur, "attn_post_norm", il); + + if (il == n_layer - 1 && inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids); + } + + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA); + cb(ffn_inp, "ffn_inp", il); + + // pre-FFN norm + cur = build_norm(ffn_inp, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + // SwiGLU dense FFN + cur = build_ffn(cur, + model.layers[il].ffn_up, NULL, NULL, + model.layers[il].ffn_gate, NULL, NULL, + model.layers[il].ffn_down, NULL, NULL, + NULL, + LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(cur, "ffn_out", il); + + cur = ggml_rms_norm(ctx0, cur, post_norm_eps); + cur = ggml_mul(ctx0, cur, model.layers[il].ffn_post_norm); + cb(cur, "ffn_post_norm", il); + + cur = ggml_add(ctx0, cur, ffn_inp); + cur = build_cvec(cur, il); + cb(cur, "l_out", il); + + inpL = cur; + } + + cur = inpL; + + // final norm + cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1); + cb(cur, "result_norm", -1); + res->t_embd = cur; + + // lm_head, followed by output multiplier + cur = build_lora_mm(model.output, cur, model.output_s); + cur = ggml_scale(ctx0, cur, hparams.f_logit_scale); + + // Final logit tanh softcap (from gemma3.cpp). + if (hparams.f_final_logit_softcapping) { + cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_final_logit_softcapping); + cur = ggml_tanh(ctx0, cur); + cur = ggml_scale(ctx0, cur, hparams.f_final_logit_softcapping); + } + + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} + +std::unique_ptr llama_model_muse_glimmer::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*this, params); +} diff --git a/src/models/nemotron-h-moe.cpp b/src/models/nemotron-h-moe.cpp index a59cc6c9f..4d03f49e0 100644 --- a/src/models/nemotron-h-moe.cpp +++ b/src/models/nemotron-h-moe.cpp @@ -1,6 +1,156 @@ #include "models.h" std::unique_ptr llama_model_nemotron_h_moe::build_arch_graph(const llm_graph_params & params) const { + if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) { + return std::make_unique(*this, params); + } return std::make_unique(*this, params); } +// MTP draft head for Nemotron-H MoE +llama_model_nemotron_h_moe::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) + : llm_graph_context(params) { + GGML_ASSERT(hparams.n_layer_nextn == 1 && "NEMOTRON_H_MOE MTP currently supports a single MTP block"); + + const int64_t n_embd_head = hparams.n_embd_head_v(); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); + + const int il = hparams.n_layer(); + const auto & layer = model.layers[il]; + + GGML_ASSERT(layer.nextn.eh_proj && layer.nextn.enorm && layer.nextn.hnorm); + GGML_ASSERT(layer.ffn_gate_inp); + + // token embedding weights + ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd; + GGML_ASSERT(tok_embd_w != nullptr && "NEMOTRON_H_MOE MTP requires token embeddings"); + + auto inp = std::make_unique(hparams.n_embd); + + inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); + ggml_set_input(inp->tokens); + + inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens); + ggml_set_input(inp->embd); + + ggml_tensor * tok_embd; + if (ubatch.token) { + tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens); + } else { + tok_embd = inp->embd; + } + cb(tok_embd, "mtp_tok_embd", il); + + inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens); + ggml_set_input(inp->h); + ggml_set_name(inp->h, "mtp_h_input"); + + ggml_tensor * h_embd = inp->h; + + res->add_input(std::move(inp)); + + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + // attention fills KV over all tokens, but the MoE is position-wise: gather output rows before + // it to save FFN compute (unless unmasked embeddings_nextn needs the full-length hidden state) + const bool emit_h_nextn = cparams.embeddings_nextn; + const bool crop_before_ffn = inp_out_ids && (!emit_h_nextn || cparams.embeddings_nextn_masked); + + auto * inp_attn = build_attn_inp_kv(); + + ggml_tensor * h_norm = build_norm(h_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il); + cb(h_norm, "mtp_hnorm", il); + ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il); + cb(e_norm, "mtp_enorm", il); + + ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, /*dim=*/ 0); + cb(concat, "mtp_concat", il); + + ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s); + cb(cur, "mtp_eh_proj", il); + + // dense NoPE attention sub-layer (mtp.layers.0) + ggml_tensor * inpSA = cur; + cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "mtp_attn_norm", il); + + { + auto [Qcur, Kcur, Vcur] = build_qkv(layer, cur, n_embd_head, hparams.n_head(il), hparams.n_head_kv(il), il); + const float kq_scale = hparams.f_attention_scale == 0.0f + ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale; + cur = build_attn(inp_attn, layer.wo, layer.wo_b, layer.wo_s, + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); + cb(cur, "mtp_attn_out", il); + } + + cur = ggml_add(ctx0, cur, inpSA); + cb(cur, "mtp_attn_residual", il); + + // gather the output rows here so the MoE FFN below only runs on the positions we keep + if (crop_before_ffn) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + + // MoE FFN sub-layer (mtp.layers.1) + ggml_tensor * ffn_residual = cur; + cur = build_norm(cur, layer.attn_post_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "mtp_attn_post_norm", il); + + { + ggml_tensor * router_logits = build_lora_mm(layer.ffn_gate_inp, cur); + cb(router_logits, "mtp_ffn_moe_logits", il); + + ggml_tensor * moe_out = + build_moe_ffn(cur, + layer.ffn_gate_inp, + layer.ffn_up_exps, + nullptr, // no gate + layer.ffn_down_exps, + layer.ffn_exp_probs_b, + n_expert, n_expert_used, + LLM_FFN_RELU_SQR, hparams.expert_weights_norm, + hparams.expert_weights_scale, + LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID, + il, + router_logits, nullptr, + layer.ffn_up_exps_s, + nullptr, // no gate + layer.ffn_down_exps_s); + cb(moe_out, "mtp_ffn_moe_out", il); + + ggml_tensor * ffn_shexp = build_ffn(cur, + layer.ffn_up_shexp, NULL, layer.ffn_up_shexp_s, + NULL, NULL, NULL, + layer.ffn_down_shexp, NULL, layer.ffn_down_shexp_s, + NULL, + LLM_FFN_RELU_SQR, LLM_FFN_PAR, il); + cb(ffn_shexp, "mtp_ffn_shexp", il); + + cur = ggml_add(ctx0, moe_out, ffn_shexp); + cb(cur, "mtp_ffn_out", il); + } + + cur = ggml_add(ctx0, cur, ffn_residual); + cb(cur, "mtp_post_ffn", il); + + // final head norm: the MTP head has its own LayerNorm + GGML_ASSERT(layer.nextn.shared_head_norm && "NEMOTRON_H_MOE MTP: missing final head norm"); + cur = build_norm(cur, layer.nextn.shared_head_norm, nullptr, LLM_NORM, -1); + + cb(cur, "h_nextn", -1); + res->t_h_nextn = cur; + + if (!crop_before_ffn && inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + + // LM head + ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output; + ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s; + GGML_ASSERT(head_w != nullptr && "NEMOTRON_H_MOE MTP requires an output projection"); + cur = build_lora_mm(head_w, cur, head_s); + cb(cur, "result_output", -1); + + res->t_logits = cur; + ggml_build_forward_expand(gf, cur); +} diff --git a/src/models/nemotron-h.cpp b/src/models/nemotron-h.cpp index a45626934..cd2af3179 100644 --- a/src/models/nemotron-h.cpp +++ b/src/models/nemotron-h.cpp @@ -7,13 +7,18 @@ void llama_model_nemotron_h::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_SSM_TIME_STEP_RANK, hparams.ssm_dt_rank); ml.get_key(LLM_KV_SSM_GROUP_COUNT, hparams.ssm_n_group); + // NextN/MTP: optional draft head appended as extra trailing block(s) + ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false); + GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < n_layer_all"); + // A layer is recurrent IFF the n_head_kv value is set to 0 and - // the n_ff value is set to 0 - for (uint32_t i = 0; i < hparams.n_layer(); ++i) { - hparams.is_recr_impl[i] = (hparams.n_head_kv(i) == 0 && hparams.n_ff(i) == 0); + // the n_ff value is set to 0. Appended MTP blocks are dense (non-recurrent) + for (uint32_t i = 0; i < hparams.n_layer_all; ++i) { + hparams.is_recr_impl[i] = i < hparams.n_layer() && hparams.n_head_kv(i) == 0 && hparams.n_ff(i) == 0; } ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps); // MTP head final_layernorm ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp, false); ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false); @@ -30,9 +35,13 @@ void llama_model_nemotron_h::load_arch_hparams(llama_model_loader & ml) { } } -void llama_model_nemotron_h::load_arch_tensors(llama_model_loader &) { +void llama_model_nemotron_h::load_arch_tensors(llama_model_loader & ml) { LLAMA_LOAD_LOCALS; + const bool mtp_only = hparams.n_layer_nextn > 0 && ml.get_weight("blk.0.attn_norm.weight") == nullptr; + const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0; + const int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0; + // mamba2 Mixer SSM params // NOTE: int64_t for tensor dimensions const int64_t d_conv = hparams.ssm_d_conv; @@ -60,61 +69,94 @@ void llama_model_nemotron_h::load_arch_tensors(llama_model_loader &) { auto & layer = layers[i]; // all blocks use the attn norm - layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, trunk_flags); if (hparams.is_recr(i)) { // ssm layers - layer.ssm_in = create_tensor(tn(LLM_TENSOR_SSM_IN, "weight", i), {n_embd, d_in_proj}, 0); + layer.ssm_in = create_tensor(tn(LLM_TENSOR_SSM_IN, "weight", i), {n_embd, d_in_proj}, trunk_flags); - layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", i), {d_conv, d_inner + 2*n_group*d_state}, 0); + layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", i), {d_conv, d_inner + 2*n_group*d_state}, trunk_flags); layer.ssm_conv1d_b = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "bias", i), {d_inner + 2*n_group*d_state}, TENSOR_NOT_REQUIRED); - layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {n_ssm_head}, 0); + layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {n_ssm_head}, trunk_flags); // no "weight" suffix for these - layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {1, n_ssm_head}, 0); - layer.ssm_d = create_tensor(tn(LLM_TENSOR_SSM_D, i), {1, n_ssm_head}, 0); + layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {1, n_ssm_head}, trunk_flags); + layer.ssm_d = create_tensor(tn(LLM_TENSOR_SSM_D, i), {1, n_ssm_head}, trunk_flags); - layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {d_inner / n_group, n_group}, 0); + layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {d_inner / n_group, n_group}, trunk_flags); // out_proj - layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", i), {d_inner, n_embd}, 0); + layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", i), {d_inner, n_embd}, trunk_flags); } else if (hparams.n_ff(i) == 0) { // attention layers (with optional bias) const int64_t n_head_i = hparams.n_head(i); const int64_t n_embd_k_gqa_i = hparams.n_embd_k_gqa(i); const int64_t n_embd_v_gqa_i = hparams.n_embd_v_gqa(i); - create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head_i, n_embd_k_gqa_i, n_embd_v_gqa_i, 0); - layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head_i, n_embd}, 0); + create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head_i, n_embd_k_gqa_i, n_embd_v_gqa_i, trunk_flags); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head_i, n_embd}, trunk_flags); layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "bias", i), {n_embd}, TENSOR_NOT_REQUIRED); } else { if (n_expert != 0) { const int64_t n_ff_exp = hparams.n_ff_exp ? hparams.n_ff_exp : n_ff / n_expert_used; const int64_t n_ff_shexp = hparams.n_ff_shexp; - layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), { n_embd, n_expert}, 0); - layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert }, 0); + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), { n_embd, n_expert}, trunk_flags); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert }, trunk_flags); // MoE branch layer.ffn_latent_down = create_tensor(tn(LLM_TENSOR_FFN_LATENT_DOWN, "weight", i), {n_embd, moe_n_embd}, TENSOR_NOT_REQUIRED); layer.ffn_latent_up = create_tensor(tn(LLM_TENSOR_FFN_LATENT_UP, "weight", i), {moe_n_embd, n_embd}, TENSOR_NOT_REQUIRED); - layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, moe_n_embd, n_expert}, 0); - layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {moe_n_embd, n_ff_exp, n_expert}, 0); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, moe_n_embd, n_expert}, trunk_flags); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {moe_n_embd, n_ff_exp, n_expert}, trunk_flags); // Shared expert branch - layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, 0); - layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, 0); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, trunk_flags); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, trunk_flags); } else { // mlp layers - layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { hparams.n_ff(i), n_embd}, 0); - layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, hparams.n_ff(i)}, 0); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { hparams.n_ff(i), n_embd}, trunk_flags); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, hparams.n_ff(i)}, trunk_flags); layer.ffn_down_b = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "bias", i), {n_embd}, TENSOR_NOT_REQUIRED); layer.ffn_up_b = create_tensor(tn(LLM_TENSOR_FFN_UP, "bias", i), {hparams.n_ff(i)}, TENSOR_NOT_REQUIRED); } } } + + // NextN/MTP draft head: each predict layer folds an attention sub-layer and a MoE + // sub-layer into a single trailing block + for (int i = n_layer; i < n_layer_all; ++i) { + auto & layer = layers[i]; + + const int64_t n_head_i = hparams.n_head(i); + const int64_t n_embd_k_gqa_i = hparams.n_embd_k_gqa(i); + const int64_t n_embd_v_gqa_i = hparams.n_embd_v_gqa(i); + const int64_t n_ff_exp = hparams.n_ff_exp ? hparams.n_ff_exp : n_ff / n_expert_used; + const int64_t n_ff_shexp = hparams.n_ff_shexp; + + // NextN input-fusion tensors + layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), {n_embd}, mtp_flags); + layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), {n_embd}, mtp_flags); + layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), {2*n_embd, n_embd}, mtp_flags); + layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), {n_embd}, mtp_flags); + + // attention sub-layer + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, mtp_flags); + create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head_i, n_embd_k_gqa_i, n_embd_v_gqa_i, mtp_flags); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head_i, n_embd}, mtp_flags); + layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "bias", i), {n_embd}, mtp_flags | TENSOR_NOT_REQUIRED); + + // MoE sub-layer + layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), {n_embd}, mtp_flags); + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, mtp_flags); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, mtp_flags); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, moe_n_embd, n_expert}, mtp_flags); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {moe_n_embd, n_ff_exp, n_expert}, mtp_flags); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, mtp_flags); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, mtp_flags); + } } std::unique_ptr llama_model_nemotron_h::build_arch_graph(const llm_graph_params & params) const { @@ -153,7 +195,7 @@ llama_model_nemotron_h::graph::graph(const llama_model & model, const llm_graph_ cur = build_ffn_layer(cur, model, il); } - if (il == n_layer - 1 && inp_out_ids) { + if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) { cur = ggml_get_rows(ctx0, cur, inp_out_ids); inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids); } @@ -170,6 +212,14 @@ llama_model_nemotron_h::graph::graph(const llama_model & model, const llm_graph_ cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1); + // seed for the MTP/NextN draft head + cb(cur, "h_nextn", -1); + res->t_h_nextn = cur; + + if (!cparams.embeddings_nextn_masked && inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + cb(cur, "result_norm", -1); res->t_embd = cur; diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index c167f9447..96fb62668 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -455,6 +455,7 @@ enum projector_type { PROJECTOR_TYPE_MIMO_AUDIO, PROJECTOR_TYPE_QWEN3TTS_SPKENC, PROJECTOR_TYPE_QWEN3TTS_GEN, + PROJECTOR_TYPE_MUSE_GLIMMER, PROJECTOR_TYPE_UNKNOWN, }; @@ -514,6 +515,7 @@ static std::map PROJECTOR_TYPE_NAMES = { { PROJECTOR_TYPE_PARAKEET, "parakeet"}, { PROJECTOR_TYPE_QWEN3TTS_SPKENC, "qwen3tts_spkenc"}, { PROJECTOR_TYPE_QWEN3TTS_GEN, "qwen3tts_gen"}, + { PROJECTOR_TYPE_MUSE_GLIMMER, "muse-glimmer"}, }; static projector_type clip_projector_type_from_string(const std::string & str) { diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index 7db01b576..761aabf64 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -109,6 +109,11 @@ struct clip_hparams { int32_t downsample_query_side; int32_t downsample_window_side; + // Muse Glimmer vision (per-block sparse-window pattern, learned pos-emb, patch-temporal) + // NOTE: these perhaps shouldn't have the architecture prefix + int32_t muse_glimmer_patch_temporal = 0; + int32_t muse_glimmer_sparse_factor = 0; + // audio int32_t n_mel_bins = 0; // whisper preprocessor int32_t proj_stack_factor = 0; // ultravox diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index 200fd23da..c74e840e8 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -69,6 +69,7 @@ #include "models/qwen3a.cpp" #include "models/mimovl.cpp" #include "models/mimo-audio.cpp" +#include "models/muse-glimmer.cpp" #include "models/qwen3tts-spkenc.cpp" #include "models/qwen3tts-gen.cpp" #include "models/step3vl.cpp" @@ -1007,6 +1008,10 @@ static std::unique_ptr clip_get_graph_builder(clip_ctx * ctx, const { builder = std::make_unique(ctx, img); } break; + case PROJECTOR_TYPE_MUSE_GLIMMER: + { + builder = std::make_unique(ctx, img); + } break; case PROJECTOR_TYPE_STEP3VL: { builder = std::make_unique(ctx, img); @@ -1653,6 +1658,17 @@ struct clip_model_loader { hparams.set_limit_image_tokens(8, 576); hparams.set_warmup_n_tokens(16*16); } break; + case PROJECTOR_TYPE_MUSE_GLIMMER: + { + hparams.n_merge = 2; // pixel-shuffle downsample after the ViT + hparams.image_resize_algo = RESIZE_ALGO_LANCZOS; + hparams.rope_theta = 10000.0f; + hparams.muse_glimmer_patch_temporal = 2; + hparams.muse_glimmer_sparse_factor = 4; // 3 sparse layers + 1 global, repeating + get_u32(KEY_SPATIAL_MERGE_SIZE, hparams.n_merge, false); + hparams.set_limit_image_tokens(1, 4096); + hparams.set_warmup_n_tokens(32*32); + } break; case PROJECTOR_TYPE_MIMOVL: { hparams.n_merge = 2; // spatial_merge_size @@ -2403,6 +2419,13 @@ struct clip_model_loader { model.mm_merger_fc2_w = get_tensor(string_format(TN_MM_MERGER_FC2, "weight")); model.mm_merger_fc2_b = get_tensor(string_format(TN_MM_MERGER_FC2, "bias")); } break; + case PROJECTOR_TYPE_MUSE_GLIMMER: + { + // 3-linear MLP: fc -> erf-GELU -> proj -> erf-GELU -> vision_proj (into LLM residual dim) + model.mm_0_w = get_tensor(string_format(TN_LLAVA_PROJ, 0, "weight")); + model.mm_1_w = get_tensor(string_format(TN_LLAVA_PROJ, 1, "weight")); + model.mm_2_w = get_tensor(string_format(TN_LLAVA_PROJ, 2, "weight")); + } break; case PROJECTOR_TYPE_STEP3VL: { model.mm_0_w = get_tensor(string_format(TN_LLAVA_PROJ, 0, "weight")); @@ -3831,6 +3854,7 @@ int clip_n_output_tokens_x(const clip_ctx * ctx, const clip_image_f32 * img) { case PROJECTOR_TYPE_PADDLEOCR: case PROJECTOR_TYPE_HUNYUANVL: case PROJECTOR_TYPE_YOUTUVL: + case PROJECTOR_TYPE_MUSE_GLIMMER: return (img->nx() / params.patch_size) / 2; case PROJECTOR_TYPE_STEP3VL: return img->nx() / (params.patch_size * params.n_merge); @@ -3856,6 +3880,7 @@ int clip_n_output_tokens_y(const clip_ctx * ctx, const clip_image_f32 * img) { case PROJECTOR_TYPE_PADDLEOCR: case PROJECTOR_TYPE_HUNYUANVL: case PROJECTOR_TYPE_YOUTUVL: + case PROJECTOR_TYPE_MUSE_GLIMMER: return (img->ny() / params.patch_size) / 2; case PROJECTOR_TYPE_STEP3VL: return img->ny() / (params.patch_size * params.n_merge); @@ -3934,6 +3959,7 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) { case PROJECTOR_TYPE_MINIMAX_M3: case PROJECTOR_TYPE_GLM4V: case PROJECTOR_TYPE_YOUTUVL: + case PROJECTOR_TYPE_MUSE_GLIMMER: { // dynamic size (2 conv, so double patch size) int x_patch = img->nx() / (params.patch_size * 2); @@ -4279,6 +4305,70 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { // set input per projector switch (ctx->model.proj_type) { + case PROJECTOR_TYPE_MUSE_GLIMMER: + { + const int grid_w = pos_w; // image_size_width / patch_size + const int grid_h = pos_h; // image_size_height / patch_size + const int n_tok = grid_w * grid_h; + const int pgrid = (int) std::sqrt((double) ctx->model.position_embeddings->ne[1]); // 32 + const int f = hparams.n_merge; // downsample 2 + + // pixel patchify runs inside the graph via build_inp() (ggml_conv_2d); + // pos-emb bilinear interp via resize_position_embeddings(). + + // --- sparse window grouping (pgrid x pgrid windows) --- + const int win = pgrid; + const int nwin_h = (grid_h + win - 1) / win; + const int nwin_w = (grid_w + win - 1) / win; + std::vector sp_perm; sp_perm.reserve(n_tok); + std::vector sp_slens; + for (int wy = 0; wy < nwin_h; wy++) { + for (int wx = 0; wx < nwin_w; wx++) { + int cnt = 0; + for (int hh = 0; hh < win; hh++) { + for (int ww = 0; ww < win; ww++) { + const int gy = wy * win + hh; + const int gx = wx * win + ww; + if (gy < grid_h && gx < grid_w) { sp_perm.push_back(gy * grid_w + gx); cnt++; } + } + } + if (cnt > 0) sp_slens.push_back(cnt); + } + } + std::vector rpos_w(n_tok), rpos_h(n_tok), inv_perm(n_tok); + for (int i = 0; i < n_tok; i++) { + const int orig = sp_perm[i]; + rpos_w[i] = (orig % grid_w) + 1; // 1-indexed + rpos_h[i] = (orig / grid_w) + 1; + inv_perm[orig] = i; + } + set_input_i32("muse_glimmer_sp_perm", sp_perm); + set_input_i32("muse_glimmer_inv_perm", inv_perm); + set_input_i32("muse_glimmer_pos_w", rpos_w); + set_input_i32("muse_glimmer_pos_h", rpos_h); + + // block-diagonal window mask (permuted order) + std::vector sp_mask((size_t) n_tok * n_tok, -INFINITY); + { + int off = 0; + for (int s : sp_slens) { + for (int a = 0; a < s; a++) + for (int b = 0; b < s; b++) + sp_mask[(size_t) (off + a) * n_tok + (off + b)] = 0.0f; + off += s; + } + } + set_input_f32("muse_glimmer_sp_mask", sp_mask); + + // pixel-shuffle gather (original order): f*f spatial neighbours grouped + std::vector dsp; dsp.reserve(n_tok); + for (int oy = 0; oy < grid_h / f; oy++) + for (int ox = 0; ox < grid_w / f; ox++) + for (int ry = 0; ry < f; ry++) + for (int rx = 0; rx < f; rx++) + dsp.push_back((oy * f + ry) * grid_w + (ox * f + rx)); + set_input_i32("muse_glimmer_ds_perm", dsp); + } break; case PROJECTOR_TYPE_MINICPMV: { // inspired from siglip: @@ -5455,6 +5545,8 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) { return ctx->model.mm_model_mlp_3_w->ne[1]; case PROJECTOR_TYPE_MINIMAX_M3: return ctx->model.mm_merger_fc2_b->ne[0]; + case PROJECTOR_TYPE_MUSE_GLIMMER: + return ctx->model.mm_2_w->ne[1]; case PROJECTOR_TYPE_QWEN2VL: case PROJECTOR_TYPE_QWEN25VL: case PROJECTOR_TYPE_EXAONE4_5: diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h index 4ee4eb374..519c0d019 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -365,3 +365,8 @@ private: ggml_tensor * build_newline_row(ggml_context * ctx0); ggml_tensor * append_rowwise_newlines(ggml_context * ctx0, ggml_tensor * tile_output); }; + +struct clip_graph_muse_glimmer : clip_graph { + clip_graph_muse_glimmer(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} + ggml_cgraph * build() override; +}; diff --git a/tools/mtmd/models/muse-glimmer.cpp b/tools/mtmd/models/muse-glimmer.cpp new file mode 100644 index 000000000..b201536f5 --- /dev/null +++ b/tools/mtmd/models/muse-glimmer.cpp @@ -0,0 +1,88 @@ +#include "models.h" + +// MuseGlimmer vision encoder: 50-layer ViT with 2D RoPE, sparse block-diagonal +// window attention (every 4th + last layer global), pixel-shuffle downsample, then +// adapter MLP + LLM's vision_projection. +// +// Several quantities are precomputed on host and fed as named graph inputs (filled in +// clip.cpp set_input, PROJECTOR_TYPE_MUSE_GLIMMER branch): +// muse_glimmer_pos_w/_h [n_tok] i32 : 1-indexed RoPE positions (sparse-permuted order) +// muse_glimmer_sp_perm [n_tok] i32 : window grouping permutation (applied after ln_pre) +// muse_glimmer_inv_perm [n_tok] i32 : inverse of sp_perm (applied after blocks) +// muse_glimmer_ds_perm [n_tok] i32 : pixel-shuffle gather (original order) +// muse_glimmer_sp_mask [n_tok, n_tok] f32 : block-diagonal window mask (sparse layers) +ggml_cgraph * clip_graph_muse_glimmer::build() { + const int ds = hparams.n_merge; // downsample factor (2) + const int sf = hparams.muse_glimmer_sparse_factor; // 4 + const int n_tok = n_patches; + const int n_out = (n_patches_x / ds) * (n_patches_y / ds); + const float rope_base = hparams.rope_theta; // 10000 + + auto inp_i32 = [&](const char * name, int64_t n) { + ggml_tensor * t = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n); + ggml_set_name(t, name); + ggml_set_input(t); + return t; + }; + + ggml_tensor * pos_w = inp_i32("muse_glimmer_pos_w", n_tok); + ggml_tensor * pos_h = inp_i32("muse_glimmer_pos_h", n_tok); + ggml_tensor * sp_perm = inp_i32("muse_glimmer_sp_perm", n_tok); + ggml_tensor * inv_perm = inp_i32("muse_glimmer_inv_perm", n_tok); + ggml_tensor * ds_perm = inp_i32("muse_glimmer_ds_perm", n_tok); + + ggml_tensor * sp_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_tok, n_tok); + ggml_set_name(sp_mask, "muse_glimmer_sp_mask"); + ggml_set_input(sp_mask); + + // patchify via build_inp (conv2d over raw pixels) + bilinear-resized learned pos-emb + ggml_tensor * x = build_inp(); // [n_embd, n_tok, 1] + x = ggml_add(ctx0, x, resize_position_embeddings(GGML_SCALE_MODE_BILINEAR)); + cb(x, "after_posemb", -1); + + // group patches into pgrid x pgrid windows (sparse attention order) + x = ggml_get_rows(ctx0, x, sp_perm); + cb(x, "after_sp_perm", -1); + + // per-layer mask: sparse layers get sp_mask, global layers (every sf-th and last) get none + std::vector attn_mask_layers(n_layer); + for (int il = 0; il < n_layer; ++il) { + const bool is_global = (il == n_layer - 1) || ((il + 1) % sf == 0); + attn_mask_layers[il] = is_global ? nullptr : sp_mask; + } + + // 2D RoPE: first half of head_dim uses width pos, second half uses height pos + auto add_pos = [&](ggml_tensor * cur, const clip_layer &) { + return build_rope_2d(ctx0, cur, pos_w, pos_h, rope_base, false); + }; + + build_vit_opts opts; + opts.attn_mask_layers = std::move(attn_mask_layers); + + // pre_ln, per-layer transformer, post_ln (all inside build_vit); reference uses exact (erf) GELU + x = build_vit(x, n_tok, NORM_TYPE_NORMAL, FFN_GELU_ERF, nullptr, add_pos, opts); + + // un-permute back to original grid order + x = ggml_get_rows(ctx0, x, inv_perm); + cb(x, "after_inv_perm", -1); + + // pixel-shuffle downsample: gather f*f spatial neighbors then concat channel-outer. + // out[c*(ds*ds)+s, o] = x[ds_perm gathered][o*(ds*ds)+s, c] + x = ggml_get_rows(ctx0, x, ds_perm); // [n_embd, n_tok], grouped + x = ggml_reshape_3d(ctx0, x, n_embd, ds * ds, n_out);// [c, s, o] + x = ggml_permute(ctx0, x, 1, 0, 2, 3); // [s, c, o] + x = ggml_cont(ctx0, x); + x = ggml_reshape_2d(ctx0, x, n_embd * ds * ds, n_out); // [6144, n_out] + cb(x, "encoder_out", -1); + + // adapter (6144->4096->4096, exact GELU each) + LLM vision_projection (4096->6656) + x = build_mm(model.mm_0_w, x); + x = ggml_gelu_erf(ctx0, x); + x = build_mm(model.mm_1_w, x); + x = ggml_gelu_erf(ctx0, x); + x = build_mm(model.mm_2_w, x); // [6656, n_out] + cb(x, "projected", -1); + + ggml_build_forward_expand(gf, x); + return gf; +} diff --git a/tools/mtmd/mtmd-image.cpp b/tools/mtmd/mtmd-image.cpp index 073d83d45..813fe493f 100644 --- a/tools/mtmd/mtmd-image.cpp +++ b/tools/mtmd/mtmd-image.cpp @@ -1615,3 +1615,65 @@ mtmd_image_preproc_out mtmd_image_preprocessor_granite::preprocess(const clip_im } return output; } + +// +// mtmd_image_preprocessor_muse_glimmer +// + +// Replicates transformers' get_aspect_ratio_preserving_size +static clip_image_size muse_glimmer_grid_size(int img_w, int img_h, int patch_hw, int max_tokens) { + double i_nph = (double) img_h / patch_hw; + double i_npw = (double) img_w / patch_hw; + const double ratio = i_nph > 0.0 ? i_npw / i_nph : 1.0; + if (i_nph * i_npw > (double) max_tokens) { + i_nph = std::sqrt((double) max_tokens / ratio); + i_npw = i_nph * ratio; + } + const int hs[2] = { (int) std::floor(i_nph), (int) std::ceil(i_nph) }; + const int ws[2] = { (int) std::floor(i_npw), (int) std::ceil(i_npw) }; + const double target_ar = (double) img_h / (double) img_w; + int best_nph = -1; + int best_npw = -1; + double best_d = 0.0; + for (int a = 0; a < 2; ++a) { + for (int b = 0; b < 2; ++b) { + const int nph = hs[a]; + const int npw = ws[b]; + if (nph < 1 || npw < 1 || nph * npw > max_tokens) { + continue; + } + const double d = std::fabs((double) nph / (double) npw - target_ar); + const int n_tokens = nph * npw; + const int best_n_tokens = best_nph * best_npw; + if (best_nph < 0 || d < best_d || (d == best_d && n_tokens > best_n_tokens)) { + best_nph = nph; + best_npw = npw; + best_d = d; + } + } + } + if (best_nph < 0) { // no candidate fit under the cap: round and clamp + best_nph = std::max(1, (int) std::lround(i_nph)); + best_npw = std::max(1, (int) std::lround(i_npw)); + } + return clip_image_size{ best_npw * patch_hw, best_nph * patch_hw }; +} + +mtmd_image_preproc_out mtmd_image_preprocessor_muse_glimmer::preprocess(const clip_image_u8 & img) { + const int patch_hw = hparams.patch_size * hparams.n_merge; + const int patch_area = hparams.patch_size * hparams.patch_size * hparams.n_merge * hparams.n_merge; + GGML_ASSERT(patch_area > 0 && hparams.image_max_pixels > 0); + const int max_tokens = hparams.image_max_pixels / patch_area; + + const clip_image_size original_size = img.get_size(); + const clip_image_size target_size = muse_glimmer_grid_size( + original_size.width, original_size.height, patch_hw, max_tokens); + + // PIL resizes directly to (target_w, target_h) -- a stretch, no padding. + clip_image_u8 resized_image; + img_tool::resize(img, resized_image, target_size, hparams.image_resize_algo, PAD_NONE); + + mtmd_image_preproc_out output; + output.append(hparams, resized_image, true); + return output; +} diff --git a/tools/mtmd/mtmd-image.h b/tools/mtmd/mtmd-image.h index ecb203f76..0669aa112 100644 --- a/tools/mtmd/mtmd-image.h +++ b/tools/mtmd/mtmd-image.h @@ -230,3 +230,9 @@ struct mtmd_image_preprocessor_granite : mtmd_image_preprocessor_llava_uhd { mtmd_image_preprocessor_granite(const clip_ctx * ctx) : mtmd_image_preprocessor_llava_uhd(ctx) {} mtmd_image_preproc_out preprocess(const clip_image_u8 & img) override; }; + +// pick the patch grid closest to the input aspect ratio under the per-image token cap, stretch-resize. +struct mtmd_image_preprocessor_muse_glimmer : mtmd_image_preprocessor { + mtmd_image_preprocessor_muse_glimmer(const clip_ctx * ctx) : mtmd_image_preprocessor(ctx) {} + mtmd_image_preproc_out preprocess(const clip_image_u8 & img) override; +}; diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 6ef4a9d3a..82b73d5cd 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -699,6 +699,12 @@ struct mtmd_context { img_end = "]<]end of image[>["; image_preproc = std::make_unique(ctx_v); } break; + case PROJECTOR_TYPE_MUSE_GLIMMER: + { + img_beg = "<|image_start|>"; + img_end = "<|image_end|>"; + image_preproc = std::make_unique(ctx_v); + } break; case PROJECTOR_TYPE_YOUTUVL: { // <|vision_start|> ... (image embeddings) ... <|vision_end|> diff --git a/tools/server/README-dev.md b/tools/server/README-dev.md index 31408f426..613017acf 100644 --- a/tools/server/README-dev.md +++ b/tools/server/README-dev.md @@ -201,7 +201,7 @@ Invoke a tool call, request body is a JSON object with: Headers: - `x-tool-cwd`: optional; if set, use as the CWD for tool; this is not part of tool's params because it's meant to be set by the runtime, not the LLM itself -- `x-tool-runtime`: optional; if set, run the tool inside this isolate instead of on the host. Only `docker-container:` is supported for now, using an already-running container +- `x-tool-runtime`: optional; if set, run the tool inside this isolate instead of on the host. Either `docker-container:` or `podman-container:`, using an already-running container, or `ssh:`, running the tool on a remote host Returns JSON object. There are two response formats (MCP tools use the same two formats: their result content is concatenated into `plain_text_response`, and RPC or tool errors are surfaced as the `error` string): diff --git a/tools/server/server-tools.cpp b/tools/server/server-tools.cpp index 27c663fdd..a4c1059ff 100644 --- a/tools/server/server-tools.cpp +++ b/tools/server/server-tools.cpp @@ -10,6 +10,7 @@ #include #include #include +#include #include #include #include @@ -25,6 +26,11 @@ # define NOMINMAX # endif # include +# include +# include +#else +# include +# include #endif namespace fs = std::filesystem; @@ -176,7 +182,7 @@ public: const std::function & on_chunk = nullptr) const = 0; }; -// shared subprocess execution helper, used by both the local and the docker-backed tools_io implementations. +// shared subprocess execution helper, used by both the local and the isolate-backed tools_io implementations. // combine_stderr=false when the raw stdout bytes must not be tainted by stderr, e.g. reading file contents. static tools_io::exec_result run_subprocess( const std::vector & args, @@ -184,7 +190,8 @@ static tools_io::exec_result run_subprocess( int timeout_secs, const std::function & on_chunk, bool combine_stderr, - const std::string & cwd = "") { + const std::string & cwd = "", + const std::string * stdin_data = nullptr) { tools_io::exec_result res; common_subproc proc; @@ -216,26 +223,59 @@ static tools_io::exec_result run_subprocess( } }); + // write stdin before reading stdout, the child drains stdin as it goes + // always close stdin, a transport client waits forever if its stdin pipe stays open + if (FILE * in = proc.stdin_file()) { + if (stdin_data != nullptr && !stdin_data->empty()) { +#if defined(_WIN32) + // pipe fds default to CRT text mode: binary keeps the bytes untranslated + _setmode(_fileno(in), _O_BINARY); +#endif + // a short write is not an error by itself, the exit code below decides + fwrite(stdin_data->data(), 1, stdin_data->size(), in); + } + fflush(in); + } + proc.close_stdin(); + FILE * f = proc.stdout_file(); std::string output; bool truncated = false; if (f) { +#if defined(_WIN32) + // pipe fds default to CRT text mode: binary keeps the bytes untranslated + _setmode(_fileno(f), _O_BINARY); +#endif + // read raw bytes, not lines: the output can hold NUL and must arrive as soon as it is ready + // keep draining past the size cap, else the child blocks on a full pipe char buf[4096]; - while (fgets(buf, sizeof(buf), f) != nullptr) { - if (!truncated) { - size_t len = strlen(buf); - if (output.size() + len <= max_output) { - output.append(buf, len); - if (on_chunk && !on_chunk(console_output_to_utf8(std::string(buf, len)))) { - proc.terminate(); - break; - } - } else { - size_t remaining = max_output - output.size(); - output.append(buf, remaining); - if (on_chunk && remaining > 0) on_chunk(console_output_to_utf8(std::string(buf, remaining))); - truncated = true; + for (;;) { +#if defined(_WIN32) + const int n = _read(_fileno(f), buf, (unsigned) sizeof(buf)); +#else + ssize_t n = read(fileno(f), buf, sizeof(buf)); + while (n < 0 && errno == EINTR) { + n = read(fileno(f), buf, sizeof(buf)); + } +#endif + if (n <= 0) { + break; + } + if (truncated) { + continue; + } + const size_t len = (size_t) n; + if (output.size() + len <= max_output) { + output.append(buf, len); + if (on_chunk && !on_chunk(console_output_to_utf8(std::string(buf, len)))) { + proc.terminate(); + break; } + } else { + size_t remaining = max_output - output.size(); + output.append(buf, remaining); + if (on_chunk && remaining > 0) on_chunk(console_output_to_utf8(std::string(buf, remaining))); + truncated = true; } } } @@ -473,7 +513,7 @@ private: } }; -// timeout for auxiliary isolate calls (stat/mkdir/ls/cp helpers); exec_shell_command uses its own +// timeout for auxiliary isolate calls (stat/mkdir/ls helpers); exec_shell_command uses its own // caller-controlled timeout instead, enforced separately in run() static constexpr int SERVER_TOOL_ISOLATE_EXEC_TIMEOUT = 15; // seconds static constexpr size_t SERVER_TOOL_ISOLATE_READ_FILE_MAX_SIZE = 64 * 1024 * 1024; // 64 MB @@ -524,33 +564,12 @@ public: } bool write_file(const std::string & path, const std::string & content) const override { - std::string abs_path = resolve(path); - - std::error_code ec; - fs::path tmp_dir = fs::temp_directory_path(ec); - if (ec) return false; - - static std::atomic tmp_counter{0}; - fs::path tmp = tmp_dir / string_format( - "llama-tools-io-isolate-%zu-%llu.tmp", - std::hash{}(std::this_thread::get_id()), - (unsigned long long) tmp_counter.fetch_add(1)); - - { - std::ofstream f(tmp, std::ios::binary); - if (!f) return false; - f << content; - if (!f) return false; - } - - bool ok = shell_run({"sh", "-c", "mkdir -p \"$(dirname \"$1\")\"", "_", abs_path}); - if (ok) { - ok = upload(tmp.string(), abs_path); - } - - std::error_code rm_ec; - fs::remove(tmp, rm_ec); - return ok; + // the content travels on stdin: no argv for the far side to re-parse, no temp file on the host + auto res = run_subprocess( + build_argv({"sh", "-c", "mkdir -p \"$(dirname \"$1\")\" && cat > \"$1\"", "_", resolve(path)}, + /*needs_stdin=*/true), + 4096, SERVER_TOOL_ISOLATE_EXEC_TIMEOUT, nullptr, true, "", &content); + return res.exit_code == 0 && !res.timed_out; } list_result list_entries(const std::string & base, int max_depth, list_kind kind) const override { @@ -612,9 +631,6 @@ protected: // a transport that re-parses its args in a remote shell (ssh) must join `inner` with shell_quote_join() virtual std::vector build_argv(const std::vector & inner, bool needs_stdin) const = 0; - // copy a host file into the isolate, `isolate_path` is absolute and its parent already exists - virtual bool upload(const std::string & host_path, const std::string & isolate_path) const = 0; - // quote `argv` into a single string that a POSIX shell re-parses into exactly `argv` static std::string shell_quote_join(const std::vector & argv) { std::string out; @@ -634,7 +650,7 @@ protected: private: std::string cwd; - // set the working directory in the command itself, docker's `-w` has no equivalent on every transport + // set the working directory in the command itself, no `-w` equivalent exists on every transport // auxiliary calls do not need this, they use the absolute paths from resolve() std::vector with_cwd(const std::vector & inner) const { if (cwd.empty()) { @@ -697,15 +713,16 @@ private: } }; -// an already-running docker container, driven through `docker exec` and `docker cp` -class tools_io_docker : public tools_io_isolate { +// an already-running container, driven through ` exec` +// docker and podman take the same verbs and the same argument order, so one class drives both +class tools_io_container : public tools_io_isolate { public: - tools_io_docker(std::string container_id, std::string cwd = "") - : tools_io_isolate(std::move(cwd)), container_id(std::move(container_id)) {} + tools_io_container(std::string bin, std::string container_id, std::string cwd = "") + : tools_io_isolate(std::move(cwd)), bin(std::move(bin)), container_id(std::move(container_id)) {} protected: std::vector build_argv(const std::vector & inner, bool needs_stdin) const override { - std::vector argv = {"docker", "exec"}; + std::vector argv = {bin, "exec"}; if (needs_stdin) { argv.push_back("-i"); } @@ -714,30 +731,118 @@ protected: return argv; } - bool upload(const std::string & host_path, const std::string & isolate_path) const override { - auto res = run_subprocess( - {"docker", "cp", host_path, container_id + ":" + isolate_path}, - 4096, SERVER_TOOL_ISOLATE_EXEC_TIMEOUT, nullptr, true); - return res.exit_code == 0 && !res.timed_out; - } - private: + std::string bin; std::string container_id; }; -// runtime spec used by --tools-runtime and the x-tool-runtime header -// this is the only scheme for now, ssh: and podman: can be added next to it -static const std::string SERVER_TOOL_RUNTIME_DOCKER_CONTAINER = "docker-container:"; +// a remote host reached over ssh +// this is remoting, not isolation: the tools can do anything the target account can do +class tools_io_ssh : public tools_io_isolate { +public: + tools_io_ssh(std::string target, std::string cwd = "") + : tools_io_isolate(std::move(cwd)), target(std::move(target)) {} + + // the target can come from a client header, and ssh reads options from its argv + // a target starting with '-' would become one, e.g. -oProxyCommand= runs on the host + static bool is_valid_target(const std::string & target) { + if (target.empty() || target[0] == '-') { + return false; + } + return std::all_of(target.begin(), target.end(), [](unsigned char c) { + return std::isalnum(c) || c == '.' || c == '-' || c == '_' || c == '@'; + }); + } + +protected: + std::vector build_argv(const std::vector & inner, bool needs_stdin) const override { + // the remote shell re-parses the command line, so `inner` travels as one quoted word + std::vector argv = ssh_argv(); + if (!needs_stdin) { + argv.push_back("-n"); + } + argv.push_back(target); + argv.push_back(shell_quote_join(inner)); + return argv; + } + +private: + std::string target; + + // there is no console here, so a prompt would hang the tool call + // key-based auth only, and the admin must trust the host key beforehand + static std::vector ssh_argv() { + return { + "ssh", + "-o", "BatchMode=yes", + "-o", "PasswordAuthentication=no", + "-o", "KbdInteractiveAuthentication=no", + "-o", "StrictHostKeyChecking=yes", + }; + } +}; + +// ":" spawns a container and owns it, "-container:" attaches to one +struct container_runtime_spec { + std::string bin; + std::string arg; // image name when spawning, container id when attaching + bool attach = false; + + static bool parse(const std::string & spec, container_runtime_spec & out) { + // docker and podman take the same verbs, hence a single implementation + static const char * engines[] = {"docker", "podman"}; + for (const char * bin : engines) { + const std::string attach_prefix = std::string(bin) + "-container:"; + if (spec.rfind(attach_prefix, 0) == 0) { + out = {bin, spec.substr(attach_prefix.size()), true}; + return true; + } + const std::string spawn_prefix = std::string(bin) + ":"; + if (spec.rfind(spawn_prefix, 0) == 0) { + out = {bin, spec.substr(spawn_prefix.size()), false}; + return true; + } + } + return false; + } + + // same risk as the ssh target: an id starting with '-' would become an engine option, + // e.g. --privileged + static bool is_valid_id(const std::string & id) { + if (id.empty() || !std::isalnum((unsigned char) id[0])) { + return false; + } + return std::all_of(id.begin(), id.end(), [](unsigned char c) { + return std::isalnum(c) || c == '.' || c == '-' || c == '_'; + }); + } +}; -// an empty runtime runs the tools on the host static std::unique_ptr make_tools_io(const json & params) { std::string cwd = json_value(params, "cwd", std::string()); std::string runtime = json_value(params, "runtime", std::string()); if (runtime.empty()) { + // an empty runtime runs the tools on the host return std::make_unique(cwd); } - if (runtime.rfind(SERVER_TOOL_RUNTIME_DOCKER_CONTAINER, 0) == 0) { - return std::make_unique(runtime.substr(SERVER_TOOL_RUNTIME_DOCKER_CONTAINER.size()), cwd); + container_runtime_spec container; + if (container_runtime_spec::parse(runtime, container)) { + // spawning belongs to the runtime that owns the container, a tool call only attaches + if (!container.attach) { + throw std::runtime_error("tool runtime must name a running container: " + runtime); + } + if (!container_runtime_spec::is_valid_id(container.arg)) { + throw std::runtime_error("invalid container id: " + container.arg); + } + return std::make_unique(container.bin, container.arg, cwd); + } + const std::string ssh_prefix = "ssh:"; + if (runtime.rfind(ssh_prefix, 0) == 0) { + std::string target = runtime.substr(ssh_prefix.size()); + if (!tools_io_ssh::is_valid_target(target)) { + throw std::runtime_error("invalid ssh target: " + target); + } + return std::make_unique(target, cwd); } // do not fall back to the host, the caller asked for an isolate throw std::runtime_error("unknown tool runtime: " + runtime); @@ -1769,81 +1874,82 @@ struct server_mcp_tool : server_tool { } }; -// owns the docker container used as the sandboxed runtime for tool invocations, as configured by -// --tools-runtime. "spawned" mode starts and stops the container itself; "existing" mode just reuses -// a container id the user already has running and never stops it. -struct server_tools_docker_runtime { - server_tools_docker_runtime(const server_tools_docker_runtime &) = delete; +// resolves --tools-runtime into the isolate that every tool call runs through +// spec() returns the runtime string make_tools_io() takes, and runs once per tool call +struct server_tools_runtime { + virtual ~server_tools_runtime() = default; + virtual std::string spec() = 0; +}; - explicit server_tools_docker_runtime(const std::string & spec) { - static const std::string docker_prefix = "docker:"; - if (spec.rfind(docker_prefix, 0) == 0) { - spawned = true; - image = spec.substr(docker_prefix.size()); - if (image.empty()) { - throw std::runtime_error("--tools-runtime docker: requires an image name"); - } - spawn(); - } else if (spec.rfind(SERVER_TOOL_RUNTIME_DOCKER_CONTAINER, 0) == 0) { - spawned = false; - container_id = spec.substr(SERVER_TOOL_RUNTIME_DOCKER_CONTAINER.size()); - if (container_id.empty()) { - throw std::runtime_error("--tools-runtime docker-container: requires a container id"); - } - } else { +// a target that already exists and needs no lifecycle +// the spec is validated once at startup, then passed straight through +struct server_tools_static_runtime : server_tools_runtime { + explicit server_tools_static_runtime(std::string spec) : runtime_spec(std::move(spec)) {} + std::string spec() override { return runtime_spec; } + +private: + std::string runtime_spec; +}; + +// owns the container the tools run in, as set by --tools-runtime ":" +// it is spawned here and stopped when the server exits +struct server_tools_container_runtime : server_tools_runtime { + server_tools_container_runtime(const server_tools_container_runtime &) = delete; + + explicit server_tools_container_runtime(const std::string & spec) { + container_runtime_spec parsed; + if (!container_runtime_spec::parse(spec, parsed)) { throw std::runtime_error("unknown --tools-runtime option: " + spec); } - } - ~server_tools_docker_runtime() { - if (spawned && !container_id.empty()) { - // closing stdin signals the container's shell (its pid 1) to exit; --rm then removes it - proc.close_stdin(); - proc.join(); + bin = parsed.bin; + image = parsed.arg; + if (image.empty()) { + throw std::runtime_error("--tools-runtime " + bin + ": requires an image name"); } + spawn(); } - // container id to use for the next tool call; respawns a spawned container that died on its own, - // or throws if an externally-managed one is no longer reachable - std::string get_container_id() { + ~server_tools_container_runtime() override { + // closing stdin signals the container's shell (its pid 1) to exit; --rm then removes it + proc.close_stdin(); + proc.join(); + } + + // respawns a container that died on its own, so the returned spec always names a running one + std::string spec() override { std::lock_guard lock(mutex); - if (!spawned) { - if (!is_running(container_id)) { - throw std::runtime_error(string_format( - "docker container \"%s\" is no longer running, restart it to keep using tools", - container_id.c_str())); - } - return container_id; - } - if (!proc.alive()) { - SRV_WRN("docker tools runtime container \"%s\" died, respawning\n", container_id.c_str()); + SRV_WRN("%s tools runtime container \"%s\" died, respawning\n", bin.c_str(), container_id.c_str()); spawn(); } - return container_id; + return bin + "-container:" + container_id; } private: - bool spawned = false; - std::string image; // spawned mode only + std::string bin; + std::string image; std::string container_id; - common_subproc proc; // spawned mode only: `docker run` client that keeps the container alive + common_subproc proc; // ` run` client that keeps the container alive std::mutex mutex; - // spawns "docker run --rm -i sh" and keeps its stdin open; the shell blocks reading stdin, + // spawns " run --rm -i sh" and keeps its stdin open; the shell blocks reading stdin, // so the container stays alive until we close it (see destructor) or it is killed from the outside void spawn() { + // create() writes over the handle it is given, so the previous one is released first + proc.join(); + std::error_code ec; fs::path cidfile = fs::temp_directory_path(ec) / string_format( "llama-tools-runtime-cid-%zu.tmp", std::hash{}(std::this_thread::get_id())); fs::remove(cidfile, ec); - std::vector args = {"docker", "run", "--rm", "-i", "--cidfile", cidfile.string(), image, "sh"}; + std::vector args = {bin, "run", "--rm", "-i", "--cidfile", path_to_utf8(cidfile), image, "sh"}; int options = subprocess_option_no_window | subprocess_option_inherit_environment | subprocess_option_search_user_path; if (!proc.create(args, options)) { - throw std::runtime_error("failed to spawn docker container for tools runtime (image: " + image + ")"); + throw std::runtime_error("failed to spawn " + bin + " container for tools runtime (image: " + image + ")"); } std::string cid; @@ -1855,15 +1961,10 @@ private: fs::remove(cidfile, ec); if (cid.empty()) { proc.terminate(); - throw std::runtime_error("timed out waiting for docker container to start (image: " + image + ")"); + throw std::runtime_error("timed out waiting for " + bin + " container to start (image: " + image + ")"); } container_id = cid; } - - static bool is_running(const std::string & id) { - auto res = run_subprocess({"docker", "inspect", "-f", "{{.State.Running}}", id}, 16, 5, nullptr, true); - return res.exit_code == 0 && !res.timed_out && res.output.rfind("true", 0) == 0; - } }; static server_tool & find_tool(std::vector> & tools, const std::string & name, bool require_stream) { @@ -1914,11 +2015,22 @@ static std::string get_header(const std::map & headers server_tools::server_tools() = default; server_tools::~server_tools() = default; +// the ":" form owns a container lifecycle +// anything else names an existing target, so only its spec is validated here at startup +static std::unique_ptr make_tools_runtime(const std::string & spec) { + container_runtime_spec parsed; + if (container_runtime_spec::parse(spec, parsed) && !parsed.attach) { + return std::make_unique(spec); + } + make_tools_io({{"runtime", spec}}); // nothing to own, just reject a bad spec now + return std::make_unique(spec); +} + void server_tools::setup(const std::vector & enabled_tools, server_mcp & mcp_mgr, const std::string & tools_runtime) { if (!tools_runtime.empty()) { - docker_runtime = std::make_unique(tools_runtime); + runtime = make_tools_runtime(tools_runtime); } if (!enabled_tools.empty()) { @@ -2016,11 +2128,11 @@ void server_tools::setup(const std::vector & enabled_tools, if (params.contains("runtime")) { params.erase("runtime"); } - auto runtime = get_header(req.headers, "x-tool-runtime"); - if (!runtime.empty()) { - params["runtime"] = runtime; - } else if (docker_runtime) { - params["runtime"] = SERVER_TOOL_RUNTIME_DOCKER_CONTAINER + docker_runtime->get_container_id(); + auto runtime_header = get_header(req.headers, "x-tool-runtime"); + if (!runtime_header.empty()) { + params["runtime"] = runtime_header; + } else if (runtime) { + params["runtime"] = runtime->spec(); } server_tool & tool = find_tool(tools, tool_name, stream); diff --git a/tools/server/server-tools.h b/tools/server/server-tools.h index ede303181..c4509ca80 100644 --- a/tools/server/server-tools.h +++ b/tools/server/server-tools.h @@ -31,7 +31,7 @@ struct server_tool { json to_json() const; }; -struct server_tools_docker_runtime; // impl detail, defined in server-tools.cpp +struct server_tools_runtime; // impl detail, defined in server-tools.cpp struct server_tools { std::vector> tools; @@ -40,8 +40,8 @@ struct server_tools { server_response queue_res; std::atomic res_id{0}; - // set when --tools-runtime is configured; owns the docker container used to run tools, if any - std::unique_ptr docker_runtime; + // set when --tools-runtime is configured; routes every tool call through an isolate + std::unique_ptr runtime; void setup(const std::vector & enabled_tools, server_mcp & mcp_mgr, diff --git a/tools/server/server.cpp b/tools/server/server.cpp index 1b2e6edb4..6d1aa4351 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -89,7 +89,7 @@ int llama_server(int argc, char ** argv) { std::setlocale(LC_NUMERIC, "C"); #ifndef _WIN32 - // Ignore SIGPIPE so the server does not crash if an MCP child exits while we are writing to its stdin + // Ignore SIGPIPE so the server does not crash if a child (MCP server, tools runtime) exits while we are writing to its stdin signal(SIGPIPE, SIG_IGN); #endif diff --git a/tools/server/tests/unit/test_tools_builtin.py b/tools/server/tests/unit/test_tools_builtin.py index 7da569d99..a69052c6d 100755 --- a/tools/server/tests/unit/test_tools_builtin.py +++ b/tools/server/tests/unit/test_tools_builtin.py @@ -14,7 +14,7 @@ PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".. GREP_MARKER = "llama_cpp_test_tools_builtin_marker_grep_search" # image the container runtime tests run their shell in -DOCKER_IMAGE = "busybox" +CONTAINER_IMAGE = "busybox" @pytest.fixture(autouse=True) @@ -151,54 +151,59 @@ def test_tools_builtin_cwd_header(): os.remove(marker_path) -def _docker_unavailable_reason() -> str | None: - """None if docker can run the image these tests use, otherwise the reason it can't.""" - docker_bin = shutil.which("docker") - if docker_bin is None: - return "docker is not installed" +def _container_engine_unavailable_reason(engine: str) -> str | None: + """None if `engine` can run the image these tests use, otherwise the reason it can't.""" + engine_bin = shutil.which(engine) + if engine_bin is None: + return f"{engine} is not installed" try: - # a daemon that answers `docker info` still cannot run a linux image when it serves - # windows containers, so probe the image itself, which also pulls it before the tests - subprocess.run([docker_bin, "run", "--rm", DOCKER_IMAGE, "true"], capture_output=True, timeout=60, check=True) + # a daemon that answers `info` still cannot run a linux image when it serves windows + # containers, so probe the image itself, which also pulls it before the tests + subprocess.run([engine_bin, "run", "--rm", CONTAINER_IMAGE, "true"], capture_output=True, timeout=60, check=True) except Exception as e: - return f"docker cannot run {DOCKER_IMAGE}: {e}" + return f"{engine} cannot run {CONTAINER_IMAGE}: {e}" return None -@pytest.fixture -def docker_container(): - reason = _docker_unavailable_reason() +@pytest.fixture(params=["docker", "podman"]) +def container_engine(request): + engine = request.param + reason = _container_engine_unavailable_reason(engine) if reason is not None: pytest.skip(reason) # ty: ignore[too-many-positional-arguments, invalid-argument-type] + return engine + +@pytest.fixture +def container_id(container_engine: str): proc = subprocess.run( - ["docker", "run", "-d", "--rm", DOCKER_IMAGE, "sleep", "300"], + [container_engine, "run", "-d", "--rm", CONTAINER_IMAGE, "sleep", "300"], capture_output=True, text=True, ) if proc.returncode != 0: - pytest.skip(f"failed to start docker container: {proc.stderr.strip()}") # ty: ignore[too-many-positional-arguments, invalid-argument-type] + pytest.skip(f"failed to start {container_engine} container: {proc.stderr.strip()}") # ty: ignore[too-many-positional-arguments, invalid-argument-type] - container_id = proc.stdout.strip() + cid = proc.stdout.strip() try: - yield container_id + yield cid finally: - subprocess.run(["docker", "rm", "-f", container_id], capture_output=True) + subprocess.run([container_engine, "rm", "-f", cid], capture_output=True) -def test_tools_builtin_runtime_header(docker_container: str): +def test_tools_builtin_runtime_header(container_engine: str, container_id: str): global server server.start() - headers = {"x-tool-runtime": f"docker-container:{docker_container}", "x-tool-cwd": "/tmp"} + headers = {"x-tool-runtime": f"{container_engine}-container:{container_id}", "x-tool-cwd": "/tmp"} - write_res = call_tool("write_file", {"path": "test.log", "content": "hello docker\n"}, headers=headers) + write_res = call_tool("write_file", {"path": "test.log", "content": "hello container\n"}, headers=headers) assert write_res["result"] == "file written successfully" read_res = call_tool("read_file", {"path": "test.log"}, headers=headers) - assert read_res["plain_text_response"] == "hello docker\n" + assert read_res["plain_text_response"] == "hello container\n" exec_res = call_tool("exec_shell_command", {"command": "cat test.log"}, headers=headers) - assert "hello docker" in exec_res["plain_text_response"] + assert "hello container" in exec_res["plain_text_response"] def test_tools_builtin_runtime_header_unknown_scheme(): @@ -208,18 +213,46 @@ def test_tools_builtin_runtime_header_unknown_scheme(): # an unknown runtime must fail, never silently fall back to running on the host res = server.make_request("POST", "/tools", data={"tool": "exec_shell_command", "params": {"command": "echo hi"}}, - headers={"x-tool-runtime": "ssh:example.com"}) + headers={"x-tool-runtime": "fake:does-not-exist"}) assert res.status_code == 500, res.body assert "unknown tool runtime" in str(res.body) +def test_tools_builtin_runtime_header_rejects_ssh_option_injection(): + global server + server.start() + + # ssh reads options from its argv, so a target starting with '-' must be rejected + res = server.make_request("POST", "/tools", + data={"tool": "exec_shell_command", "params": {"command": "echo hi"}}, + headers={"x-tool-runtime": "ssh:-oProxyCommand=touch /tmp/pwned"}) + assert res.status_code == 500, res.body + assert "invalid ssh target" in str(res.body) + + +@pytest.mark.parametrize("engine", ["docker", "podman"]) +def test_tools_builtin_runtime_header_rejects_container_option_injection(engine: str): + global server + server.start() + + # the container id lands on the ` exec` command line, so an id that looks + # like an option must be rejected + res = server.make_request("POST", "/tools", + data={"tool": "exec_shell_command", "params": {"command": "echo hi"}}, + headers={"x-tool-runtime": f"{engine}-container:--privileged"}) + assert res.status_code == 500, res.body + assert "invalid container id" in str(res.body) + + def test_tools_builtin_docker_runtime_cleans_up_spawned_container(): - reason = _docker_unavailable_reason() + # docker-only: this reads the container hostname to get the spawned id, which only docker + # sets to the short id. podman is covered by the attach path above + reason = _container_engine_unavailable_reason("docker") if reason is not None: pytest.skip(reason) # ty: ignore[too-many-positional-arguments, invalid-argument-type] global server - server.server_tools_runtime = f"docker:{DOCKER_IMAGE}" + server.server_tools_runtime = f"docker:{CONTAINER_IMAGE}" server.start() # exec_shell_command runs inside the container spawned for --tools-runtime; docker sets diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContenteditable.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContenteditable.svelte index 98b082c89..35977dfc3 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContenteditable.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContenteditable.svelte @@ -797,7 +797,7 @@ data-placeholder={placeholder} tabindex={disabled ? -1 : 0} class={[ - 'chat-form-contenteditable text-md min-h-12 w-full whitespace-pre-wrap wrap-break-word border-0 bg-transparent p-0 leading-6 outline-none focus-visible:ring-0 focus-visible:ring-offset-0', + 'chat-form-contenteditable text-md min-h-12 w-full overflow-y-auto whitespace-pre-wrap wrap-break-word border-0 bg-transparent p-0 leading-6 outline-none focus-visible:ring-0 focus-visible:ring-offset-0', disabled && 'cursor-not-allowed' ]} style="max-height: var(--max-message-height);" diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageSystem/ChatMessageSystem.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageSystem/ChatMessageSystem.svelte index 5a1e4f213..30605ba9f 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageSystem/ChatMessageSystem.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageSystem/ChatMessageSystem.svelte @@ -164,7 +164,7 @@ ? `max-height: ${MAX_HEIGHT}px;` : 'max-height: none;'} > - {#if !currentConfig.renderContentAsRawText} + {#if currentConfig.renderUserContentAsMarkdown}
diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUserBubble.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUserBubble.svelte index a2e91b94c..f5a3ff6f9 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUserBubble.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUserBubble.svelte @@ -1,5 +1,5 @@ + + + + + {#each getMentionBadgeIconPaths(path) as d (d)} + + {/each} + + + {label} + diff --git a/tools/ui/src/lib/components/app/content/MentionText.svelte b/tools/ui/src/lib/components/app/content/MentionText.svelte new file mode 100644 index 000000000..0a4bc0eeb --- /dev/null +++ b/tools/ui/src/lib/components/app/content/MentionText.svelte @@ -0,0 +1,17 @@ + + + + +{#each segments as segment, index (index)}{#if segment.mention}{:else}{segment.text}{/if}{/each} diff --git a/tools/ui/src/lib/components/app/content/index.ts b/tools/ui/src/lib/components/app/content/index.ts index 5cfdd1b9c..9b43fe9cc 100644 --- a/tools/ui/src/lib/components/app/content/index.ts +++ b/tools/ui/src/lib/components/app/content/index.ts @@ -31,6 +31,20 @@ */ export { default as MarkdownContent } from './MarkdownContent/MarkdownContent.svelte'; +/** + * **MentionText** - Plain text with file mention badges + * + * Renders a message verbatim, turning only `[name](file://path)` links + * into the same badge chips the markdown path draws. Nothing else is + * interpreted, so pasted code keeps its `#` comments and underscores. + * + * @example + * ```svelte + * + * ``` + */ +export { default as MentionText } from './MentionText.svelte'; + /** * **SyntaxHighlightedCode** - Code syntax highlighting * diff --git a/tools/ui/src/lib/constants/mention-badge.ts b/tools/ui/src/lib/constants/mention-badge.ts index 5de827014..a9ef963c3 100644 --- a/tools/ui/src/lib/constants/mention-badge.ts +++ b/tools/ui/src/lib/constants/mention-badge.ts @@ -10,6 +10,9 @@ export const MENTION_BADGE_CLASSNAME = export const MENTION_BADGE_ICON_CLASSNAME = 'h-3 w-3 shrink-0'; +/** Regex flag that makes the mention scanner walk every link in a message instead of the first. */ +export const MENTION_LINK_SCAN_FLAGS = 'g'; + /** * SVG attributes shared by the DOM-built and hast-built badge icons. * The tokenizer applies them via `setAttribute`, the rehype plugin diff --git a/tools/ui/src/lib/constants/path-display.ts b/tools/ui/src/lib/constants/path-display.ts index 95877d741..fd1001761 100644 --- a/tools/ui/src/lib/constants/path-display.ts +++ b/tools/ui/src/lib/constants/path-display.ts @@ -12,6 +12,9 @@ import { UrlProtocol } from '$lib/enums'; export const CWD_CHANGED_PREFIX = 'Set working directory to '; export const CWD_CLEARED_TEXT = 'Working directory cleared'; +/** Trailing separator that marks a path as a directory. */ +export const DIRECTORY_PATH_SUFFIX = '/'; + export const HOME_TILDE = '~'; export const HOME_TILDE_PREFIX = '~/'; // tilde plus path separator diff --git a/tools/ui/src/lib/constants/settings-keys.ts b/tools/ui/src/lib/constants/settings-keys.ts index 8388f2b75..b53d11048 100644 --- a/tools/ui/src/lib/constants/settings-keys.ts +++ b/tools/ui/src/lib/constants/settings-keys.ts @@ -41,7 +41,8 @@ export const SETTINGS_KEYS = { // Performance PRE_ENCODE_CONVERSATION: 'preEncodeConversation', PRESENCE_PENALTY: 'presence_penalty', - RENDER_CONTENT_AS_RAW_TEXT: 'renderContentAsRawText', + RENDER_THINKING_AS_MARKDOWN: 'renderThinkingAsMarkdown', + RENDER_USER_CONTENT_AS_MARKDOWN: 'renderUserContentAsMarkdown', // Penalties REPEAT_LAST_N: 'repeat_last_n', REPEAT_PENALTY: 'repeat_penalty', diff --git a/tools/ui/src/lib/constants/settings-registry.ts b/tools/ui/src/lib/constants/settings-registry.ts index 19db43963..8e1cc59f8 100644 --- a/tools/ui/src/lib/constants/settings-registry.ts +++ b/tools/ui/src/lib/constants/settings-registry.ts @@ -230,10 +230,18 @@ const SETTINGS_REGISTRY: Record = { type: SettingsFieldType.CHECKBOX }, { - defaultValue: false, - help: 'Display user, system and thinking content as plain text instead of formatted Markdown. Markdown is the default so that @-mention badges render in sent messages.', - key: SETTINGS_KEYS.RENDER_CONTENT_AS_RAW_TEXT, - label: 'Render content as raw text', + defaultValue: true, + help: 'Render user messages using markdown formatting in the chat. Turn this off to keep a message exactly as typed; @-mention badges show either way.', + key: SETTINGS_KEYS.RENDER_USER_CONTENT_AS_MARKDOWN, + label: 'Render user content as Markdown', + section: SETTINGS_SECTION_SLUGS.DISPLAY, + type: SettingsFieldType.CHECKBOX + }, + { + defaultValue: true, + help: 'Render the reasoning/thinking block content as formatted Markdown instead of plain text.', + key: SETTINGS_KEYS.RENDER_THINKING_AS_MARKDOWN, + label: 'Render thinking as Markdown', section: SETTINGS_SECTION_SLUGS.DISPLAY, type: SettingsFieldType.CHECKBOX }, diff --git a/tools/ui/src/lib/services/migration.service.ts b/tools/ui/src/lib/services/migration.service.ts index b2626c794..ffe4bd846 100644 --- a/tools/ui/src/lib/services/migration.service.ts +++ b/tools/ui/src/lib/services/migration.service.ts @@ -629,6 +629,37 @@ const configTypesMigration: Migration = { console.log(`[Migration] Config types: coerced string booleans (changed=${changed})`); } }; +const RENDER_KEYS_MIGRATION_ID = 'render-keys-unfold-v1'; +const LEGACY_RENDER_RAW_TEXT_KEY = 'renderContentAsRawText'; +const renderKeysMigration: Migration = { + description: 'Unfold the single raw text render toggle onto the per-surface render keys', + id: RENDER_KEYS_MIGRATION_ID, + + async run(): Promise { + const configRaw = localStorage.getItem(CONFIG_LOCALSTORAGE_KEY); + + if (configRaw === null) return; + + const config = JSON.parse(configRaw); + + if (!(LEGACY_RENDER_RAW_TEXT_KEY in config)) return; + + // The toggle carried user content and thinking at once and cannot say which surface + // was chosen, so it only restores the user key and thinking keeps its own default. + if (!(SETTINGS_KEYS.RENDER_USER_CONTENT_AS_MARKDOWN in config)) { + config[SETTINGS_KEYS.RENDER_USER_CONTENT_AS_MARKDOWN] = + config[LEGACY_RENDER_RAW_TEXT_KEY] !== true; + } + + // Dropped rather than preserved: the two render keys and the toggle describe the same + // surfaces, so leaving it behind would let a stale value fight the restored one. + delete config[LEGACY_RENDER_RAW_TEXT_KEY]; + localStorage.setItem(CONFIG_LOCALSTORAGE_KEY, JSON.stringify(config)); + + if (import.meta.env.DEV && import.meta.env.VITE_DEBUG) + console.log('[Migration] Render keys: unfolded the raw text toggle'); + } +}; const MCP_DEFAULT_OVERRIDES_LEGACY_KEY = `${STORAGE_APP_NAME}.mcpDefaultServerOverrides`; const MCP_DEFAULT_OVERRIDES_MERGE_MIGRATION_ID = 'mcp-default-overrides-merge-v1'; /** @@ -722,7 +753,8 @@ const migrations: Migration[] = [ customJsonKeyMigration, mcpDefaultEnabledMigration, mcpDefaultOverridesMergeMigration, - configTypesMigration + configTypesMigration, + renderKeysMigration ]; export const MigrationService = { diff --git a/tools/ui/src/lib/stores/settings.svelte.ts b/tools/ui/src/lib/stores/settings.svelte.ts index f413638ef..c8b9be7a7 100644 --- a/tools/ui/src/lib/stores/settings.svelte.ts +++ b/tools/ui/src/lib/stores/settings.svelte.ts @@ -135,32 +135,6 @@ class SettingsStore { ...savedVal }; - // Migrate the legacy render keys into `renderContentAsRawText` - // (inverted semantics: the old keys opted INTO markdown). Any - // explicit raw-text preference wins when the legacy keys disagree. - const LEGACY_MARKDOWN_KEYS = ['renderUserContentAsMarkdown', 'renderThinkingAsMarkdown']; - const LEGACY_RAW_TEXT_KEY = 'renderUserContentAsRawText'; // this branch's intermediate key - const legacyKeys = [...LEGACY_MARKDOWN_KEYS, LEGACY_RAW_TEXT_KEY].filter( - (key) => key in savedVal - ); - - if (legacyKeys.length > 0) { - if (!(SETTINGS_KEYS.RENDER_CONTENT_AS_RAW_TEXT in savedVal)) { - if (LEGACY_RAW_TEXT_KEY in savedVal) { - this.config[SETTINGS_KEYS.RENDER_CONTENT_AS_RAW_TEXT] = savedVal[LEGACY_RAW_TEXT_KEY]; - } else { - this.config[SETTINGS_KEYS.RENDER_CONTENT_AS_RAW_TEXT] = LEGACY_MARKDOWN_KEYS.filter( - (key) => key in savedVal - ).some((key) => savedVal[key] === false); - } - } - - for (const key of legacyKeys) { - delete (this.config as Record)[key]; - } - this.saveConfig(); - } - // Default sendOnEnter to false on mobile when the user has no saved preference if (!(SETTINGS_KEYS.SEND_ON_ENTER in savedVal)) { if (isMobile.current) { diff --git a/tools/ui/src/lib/utils/index.ts b/tools/ui/src/lib/utils/index.ts index 6703da373..5762e8e4b 100644 --- a/tools/ui/src/lib/utils/index.ts +++ b/tools/ui/src/lib/utils/index.ts @@ -240,6 +240,7 @@ export { MENTION_BADGE_FOLDER_ICON_PATHS, getMentionBadgeIconPaths, getMentionBadgeLabel, + splitMentionSegments, buildMentionInsertion } from './mention-badge'; diff --git a/tools/ui/src/lib/utils/mention-badge.ts b/tools/ui/src/lib/utils/mention-badge.ts index 7b76a2280..5b1b7d1fe 100644 --- a/tools/ui/src/lib/utils/mention-badge.ts +++ b/tools/ui/src/lib/utils/mention-badge.ts @@ -1,8 +1,9 @@ import { abbreviateHome, lastPathSegment } from './path-display'; -import { FILE_URI_PREFIX } from '$lib/constants'; +import { DIRECTORY_PATH_SUFFIX, FILE_URI_PREFIX } from '$lib/constants'; import { MENTION_BADGE_FILE_ICON_PATHS, - MENTION_BADGE_FOLDER_ICON_PATHS + MENTION_BADGE_FOLDER_ICON_PATHS, + MENTION_LINK_SCAN_FLAGS } from '$lib/constants/mention-badge'; import { FileMentionEntryType } from '$lib/enums'; import type { FileMentionEntry } from '$lib/types'; @@ -48,8 +49,45 @@ export function decodeFileLinkPath(path: string): string { } } +export interface MentionTextSegment { + text: string; + mention: { name: string; path: string } | null; +} + +/** + * Split raw text into plain runs and `[name](file://path)` mentions. + * The raw-text renderers walk these segments to draw badges without + * handing the message to the markdown parser, so a `#` stays a `#`. + */ +export function splitMentionSegments(value: string): MentionTextSegment[] { + const linkRe = fileMentionLinkRe(MENTION_LINK_SCAN_FLAGS); + const segments: MentionTextSegment[] = []; + + let cursor = 0; + let match: RegExpExecArray | null; + + while ((match = linkRe.exec(value)) !== null) { + if (match.index > cursor) { + segments.push({ mention: null, text: value.slice(cursor, match.index) }); + } + + segments.push({ + mention: { name: match[1], path: decodeFileLinkPath(match[2]) }, + text: match[0] + }); + + cursor = match.index + match[0].length; + } + + if (cursor < value.length) segments.push({ mention: null, text: value.slice(cursor) }); + + return segments; +} + export function getMentionBadgeIconPaths(path: string): readonly string[] { - return path.endsWith('/') ? MENTION_BADGE_FOLDER_ICON_PATHS : MENTION_BADGE_FILE_ICON_PATHS; + return path.endsWith(DIRECTORY_PATH_SUFFIX) + ? MENTION_BADGE_FOLDER_ICON_PATHS + : MENTION_BADGE_FILE_ICON_PATHS; } export function getMentionBadgeLabel( diff --git a/tools/ui/tests/client/settings-raw-text-migration.svelte.test.ts b/tools/ui/tests/client/settings-raw-text-migration.svelte.test.ts deleted file mode 100644 index 7e60ccfb9..000000000 --- a/tools/ui/tests/client/settings-raw-text-migration.svelte.test.ts +++ /dev/null @@ -1,66 +0,0 @@ -// Guards the legacy render-key migration: `renderUserContentAsMarkdown` -// and `renderThinkingAsMarkdown` (opt-INTO markdown) fold into the single -// `renderContentAsRawText` setting, with any explicit raw-text preference -// winning when the legacy keys disagree. Legacy keys are removed from the -// persisted config so they do not stay orphaned in localStorage. - -import { CONFIG_LOCALSTORAGE_KEY } from '$lib/constants/storage'; -import { config, settingsStore } from '$lib/stores/settings.svelte'; -import { beforeEach, describe, expect, it } from 'vitest'; - -function seedConfig(stored: Record) { - localStorage.setItem(CONFIG_LOCALSTORAGE_KEY, JSON.stringify(stored)); - settingsStore.initialize(); -} - -function persisted(): Record { - return JSON.parse(localStorage.getItem(CONFIG_LOCALSTORAGE_KEY) ?? '{}'); -} - -describe('renderContentAsRawText migration', () => { - beforeEach(() => { - localStorage.removeItem(CONFIG_LOCALSTORAGE_KEY); - settingsStore.initialize(); - }); - - it('maps renderUserContentAsMarkdown=false to raw text', () => { - seedConfig({ renderUserContentAsMarkdown: false }); - expect(config().renderContentAsRawText).toBe(true); - }); - - it('maps renderUserContentAsMarkdown=true to markdown', () => { - seedConfig({ renderUserContentAsMarkdown: true }); - expect(config().renderContentAsRawText).toBe(false); - }); - - it('maps renderThinkingAsMarkdown=false to raw text', () => { - seedConfig({ renderThinkingAsMarkdown: false }); - expect(config().renderContentAsRawText).toBe(true); - }); - - it('lets any explicit raw-text preference win when the legacy keys disagree', () => { - seedConfig({ renderThinkingAsMarkdown: false, renderUserContentAsMarkdown: true }); - expect(config().renderContentAsRawText).toBe(true); - }); - - it('honors the intermediate renderUserContentAsRawText key from the PR branch', () => { - seedConfig({ renderUserContentAsRawText: true }); - expect(config().renderContentAsRawText).toBe(true); - }); - - it('keeps an already-migrated value and cleans up the legacy keys', () => { - seedConfig({ renderContentAsRawText: false, renderUserContentAsMarkdown: false }); - expect(config().renderContentAsRawText).toBe(false); - - const stored = persisted(); - - expect(stored.renderUserContentAsMarkdown).toBeUndefined(); - expect(stored.renderThinkingAsMarkdown).toBeUndefined(); - expect(stored.renderUserContentAsRawText).toBeUndefined(); - }); - - it('defaults to markdown when no legacy key exists', () => { - seedConfig({}); - expect(config().renderContentAsRawText).toBe(false); - }); -}); diff --git a/tools/ui/tests/client/settings-render-keys-migration.svelte.test.ts b/tools/ui/tests/client/settings-render-keys-migration.svelte.test.ts new file mode 100644 index 000000000..380a2e74d --- /dev/null +++ b/tools/ui/tests/client/settings-render-keys-migration.svelte.test.ts @@ -0,0 +1,64 @@ +// Guards the unfolding of `renderContentAsRawText` back onto the two +// per-surface render keys. The single toggle carried user content and +// thinking at once, so only the user key is restored from it and thinking +// returns to its own default. The toggle is removed from the persisted +// config so it does not stay orphaned in localStorage. + +import { CONFIG_LOCALSTORAGE_KEY } from '$lib/constants/storage'; +import { MigrationService } from '$lib/services/migration.service'; +import { config, settingsStore } from '$lib/stores/settings.svelte'; +import { beforeEach, describe, expect, it } from 'vitest'; + +const RENDER_KEYS_MIGRATION_ID = 'render-keys-unfold-v1'; + +async function seedConfig(stored: Record) { + localStorage.setItem(CONFIG_LOCALSTORAGE_KEY, JSON.stringify(stored)); + + const migration = MigrationService.getMigrations().find((m) => m.id === RENDER_KEYS_MIGRATION_ID); + + await migration?.run(); + settingsStore.initialize(); +} + +function persisted(): Record { + return JSON.parse(localStorage.getItem(CONFIG_LOCALSTORAGE_KEY) ?? '{}'); +} + +describe('renderContentAsRawText unfolding', () => { + beforeEach(() => { + localStorage.removeItem(CONFIG_LOCALSTORAGE_KEY); + MigrationService.resetState(); + settingsStore.initialize(); + }); + + it('maps raw text to user content as plain text', async () => { + await seedConfig({ renderContentAsRawText: true }); + expect(config().renderUserContentAsMarkdown).toBe(false); + }); + + it('maps markdown to user content as markdown', async () => { + await seedConfig({ renderContentAsRawText: false }); + expect(config().renderUserContentAsMarkdown).toBe(true); + }); + + it('leaves thinking on its own default', async () => { + await seedConfig({ renderContentAsRawText: true }); + expect(config().renderThinkingAsMarkdown).toBe(true); + }); + + it('keeps an explicit user preference over the toggle', async () => { + await seedConfig({ renderContentAsRawText: true, renderUserContentAsMarkdown: true }); + expect(config().renderUserContentAsMarkdown).toBe(true); + }); + + it('drops the toggle from the persisted config', async () => { + await seedConfig({ renderContentAsRawText: true }); + expect(persisted().renderContentAsRawText).toBeUndefined(); + }); + + it('leaves both surfaces on markdown when nothing is stored', async () => { + await seedConfig({}); + expect(config().renderUserContentAsMarkdown).toBe(true); + expect(config().renderThinkingAsMarkdown).toBe(true); + }); +}); diff --git a/tools/ui/tests/unit/mention-segments.test.ts b/tools/ui/tests/unit/mention-segments.test.ts new file mode 100644 index 000000000..51bb46d1e --- /dev/null +++ b/tools/ui/tests/unit/mention-segments.test.ts @@ -0,0 +1,48 @@ +import { splitMentionSegments } from '$lib/utils/mention-badge'; +import { describe, expect, it } from 'vitest'; + +describe('splitMentionSegments', () => { + it('returns a single plain run when there is no mention', () => { + const segments = splitMentionSegments('# not a heading here'); + + expect(segments).toEqual([{ mention: null, text: '# not a heading here' }]); + }); + + it('splits text around a mention', () => { + const segments = splitMentionSegments('look at [main.c](file:///src/main.c) please'); + + expect(segments.map((segment) => segment.text)).toEqual([ + 'look at ', + '[main.c](file:///src/main.c)', + ' please' + ]); + expect(segments[1].mention).toEqual({ name: 'main.c', path: '/src/main.c' }); + }); + + it('decodes percent-encoded paths', () => { + const segments = splitMentionSegments('[a b.txt](file:///tmp/a%20b.txt)'); + + expect(segments[0].mention?.path).toBe('/tmp/a b.txt'); + }); + + it('keeps the directory marker so the folder icon is picked', () => { + const segments = splitMentionSegments('[src](file:///repo/src/)'); + + expect(segments[0].mention?.path).toBe('/repo/src/'); + }); + + it('handles adjacent mentions with no text between them', () => { + const segments = splitMentionSegments('[a](file:///a)[b](file:///b)'); + + expect(segments).toHaveLength(2); + expect(segments.every((segment) => segment.mention !== null)).toBe(true); + }); + + it('preserves the exact source when segments are joined back', () => { + const source = 'see [a](file:///a) and [b](file:///b/) done'; + + expect(splitMentionSegments(source).reduce((acc, segment) => acc + segment.text, '')).toBe( + source + ); + }); +}); diff --git a/vendor/cpp-httplib/httplib.cpp b/vendor/cpp-httplib/httplib.cpp index 30e2de896..3b687ff4f 100644 --- a/vendor/cpp-httplib/httplib.cpp +++ b/vendor/cpp-httplib/httplib.cpp @@ -385,8 +385,7 @@ void set_verify_client(ctx_t ctx, bool require); // Session management session_t create_session(ctx_t ctx, socket_t sock); void free_session(session_t session); -bool set_sni(session_t session, const char *hostname); -bool set_hostname(session_t session, const char *hostname); +bool set_sni(session_t session, const char *hostname, bool verify_hostname); // Handshake (non-blocking capable) TlsError connect(session_t session); @@ -1877,6 +1876,39 @@ int shutdown_socket(socket_t sock) noexcept { #endif } +// Half-closes the write side and drains any in-flight/queued bytes before +// the final shutdown+close. Closing with unread data in the receive queue +// (or bytes arriving after the receive side is closed) makes the stack send +// an abortive RST instead of a graceful FIN, which can make the peer see the +// response as a failed read even though it was fully written. +void drain_and_close_socket(socket_t sock) noexcept { +#ifdef _WIN32 + shutdown(sock, SD_SEND); +#else + shutdown(sock, SHUT_WR); +#endif + + char buf[CPPHTTPLIB_RECV_BUFSIZ]; + size_t total = 0; + const auto deadline = std::chrono::steady_clock::now() + + std::chrono::milliseconds(100); // bound #1 + + while (total < size_t(1024u * 1024u)) { // bound #2 + const auto remaining = + std::chrono::duration_cast( + deadline - std::chrono::steady_clock::now()) + .count(); + if (remaining <= 0) { break; } + if (select_read(sock, 0, static_cast(remaining)) <= 0) { break; } + const auto n = read_socket(sock, buf, sizeof(buf), CPPHTTPLIB_RECV_FLAGS); + if (n <= 0) { break; } + total += static_cast(n); + } + + shutdown_socket(sock); + close_socket(sock); +} + std::string escape_abstract_namespace_unix_domain(const std::string &s) { if (s.size() > 1 && s[0] == '\0') { auto ret = s; @@ -3320,42 +3352,94 @@ bool read_headers(Stream &strm, Headers &headers) { return true; } +bool parse_status_line(const char *line, std::string &version, + int &status, std::string &reason) { +#ifdef CPPHTTPLIB_ALLOW_LF_AS_LINE_TERMINATOR + thread_local const std::regex re("(HTTP/1\\.[01]) (\\d{3})(?: (.*?))?\r?\n"); +#else + thread_local const std::regex re("(HTTP/1\\.[01]) (\\d{3})(?: (.*?))?\r\n"); +#endif + + std::cmatch m; + if (!std::regex_match(line, m, re)) { return false; } + version = std::string(m[1]); + status = std::stoi(std::string(m[2])); + reason = std::string(m[3]); + return true; +} + +// Everything WebSocketClient::connect() reports about the upgrade exchange. +// status stays -1 until a status line is parsed, mirroring stream::Result. +struct WebSocketUpgradeResponse { + Error error = Error::Success; + int status = -1; + Headers headers; + std::string selected_subprotocol; +}; + bool read_websocket_upgrade_response(Stream &strm, const std::string &expected_accept, - std::string &selected_subprotocol) { + WebSocketUpgradeResponse &upgrade) { // Read status line const auto bufsiz = 2048; char buf[bufsiz]; stream_line_reader line_reader(strm, buf, bufsiz); - if (!line_reader.getline()) { return false; } + if (!line_reader.getline()) { + upgrade.error = Error::Read; + return false; + } - // Check for "HTTP/1.1 101" - auto line = std::string(line_reader.ptr(), line_reader.size()); - if (line.find("HTTP/1.1 101") == std::string::npos) { return false; } + std::string version; + std::string reason; + if (!parse_status_line(line_reader.ptr(), version, upgrade.status, reason)) { + upgrade.error = Error::WebSocketHandshake; + return false; + } - // Parse headers using existing read_headers - Headers headers; - if (!read_headers(strm, headers)) { return false; } + // Read the headers even for a rejection so the caller can see why the + // server refused the upgrade. A non-101 response may carry a body; it is + // deliberately left unread since the caller closes the socket right away. + if (!read_headers(strm, upgrade.headers)) { + upgrade.error = Error::Read; + return false; + } + + const auto &headers = upgrade.headers; + + if (upgrade.status != StatusCode::SwitchingProtocol_101) { + upgrade.error = Error::WebSocketHandshake; + return false; + } // Verify Upgrade: websocket (case-insensitive) auto upgrade_it = headers.find("Upgrade"); - if (upgrade_it == headers.end()) { return false; } - auto upgrade_val = case_ignore::to_lower(upgrade_it->second); - if (upgrade_val != "websocket") { return false; } + if (upgrade_it == headers.end() || + case_ignore::to_lower(upgrade_it->second) != "websocket") { + upgrade.error = Error::WebSocketHandshake; + return false; + } // Verify Connection header contains "Upgrade" (case-insensitive) auto connection_it = headers.find("Connection"); - if (connection_it == headers.end()) { return false; } - auto connection_val = case_ignore::to_lower(connection_it->second); - if (connection_val.find("upgrade") == std::string::npos) { return false; } + if (connection_it == headers.end() || + case_ignore::to_lower(connection_it->second).find("upgrade") == + std::string::npos) { + upgrade.error = Error::WebSocketHandshake; + return false; + } // Verify Sec-WebSocket-Accept header value auto it = headers.find("Sec-WebSocket-Accept"); - if (it == headers.end() || it->second != expected_accept) { return false; } + if (it == headers.end() || it->second != expected_accept) { + upgrade.error = Error::WebSocketHandshake; + return false; + } // Extract negotiated subprotocol auto proto_it = headers.find("Sec-WebSocket-Protocol"); - if (proto_it != headers.end()) { selected_subprotocol = proto_it->second; } + if (proto_it != headers.end()) { + upgrade.selected_subprotocol = proto_it->second; + } return true; } @@ -5105,7 +5189,7 @@ bool is_field_valid(const std::string &name, const std::string &value) { } // namespace fields bool perform_websocket_handshake(Stream &strm, Request &req, - std::string &selected_subprotocol) { + WebSocketUpgradeResponse &upgrade) { // Generate random Sec-WebSocket-Key thread_local std::mt19937 rng(std::random_device{}()); std::string key_bytes(16, '\0'); @@ -5130,20 +5214,26 @@ bool perform_websocket_handshake(Stream &strm, Request &req, // and would emit one small write per header. BufferStream bstrm; - if (write_request_line(bstrm, req.method, req.path) < 0) { return false; } + if (write_request_line(bstrm, req.method, req.path) < 0) { + upgrade.error = Error::Write; + return false; + } auto error = Error::Success; if (!check_and_write_headers(bstrm, req.headers, write_headers, error)) { + upgrade.error = error; return false; } const auto &data = bstrm.get_buffer(); - if (!write_data(strm, data.data(), data.size())) { return false; } + if (!write_data(strm, data.data(), data.size())) { + upgrade.error = Error::Write; + return false; + } // Verify 101 response and Sec-WebSocket-Accept header auto expected_accept = websocket_accept_key(client_key); - return read_websocket_upgrade_response(strm, expected_accept, - selected_subprotocol); + return read_websocket_upgrade_response(strm, expected_accept, upgrade); } bool is_ip_address(const std::string &host) { @@ -5637,50 +5727,144 @@ bool load_client_ca_config(tls::ctx_t ctx, return ret; } -bool setup_client_tls_session(const std::string &host, tls::ctx_t ctx, - tls::session_t &session, socket_t sock, - bool server_certificate_verification, - time_t timeout_sec, time_t timeout_usec) { +// The parts of session setup that only SSLClient needs, plus the handful +// WebSocketClient also exposes; everything else takes the defaults, which is +// what keeps the two clients on one implementation. +struct ClientTlsSessionOptions { + // Both SSLClient and WebSocketClient expose this independently of + // certificate verification. + bool server_hostname_verification = true; + std::function session_verifier; + // When non-null, guards session creation against concurrent use of the + // context. A WebSocketClient is not safe to use from several threads to + // begin with, so it passes nothing. + std::mutex *ctx_mutex = nullptr; +#ifdef CPPHTTPLIB_WINDOWS_AUTOMATIC_ROOT_CERTIFICATES_UPDATE + // The caller decides whether Schannel has anything to say about this + // connection; see SSLClient::initialize_ssl(). + bool windows_cert_verification = false; +#endif +}; + +// Filled in on failure for callers that report error details. +struct ClientTlsSessionError { + Error error = Error::Success; + int ssl_error = 0; + uint64_t backend_error = 0; +}; + +// Establishes a client TLS session on an already connected socket. On failure +// the session is left for the caller to free: SSLClient frees it right away, +// WebSocketClient keeps it in a member that shutdown_and_close() cleans up. +bool setup_client_tls_session( + const std::string &host, tls::ctx_t ctx, tls::session_t &session, + socket_t sock, bool server_certificate_verification, time_t timeout_sec, + time_t timeout_usec, ClientTlsSessionError *out_error = nullptr, + const ClientTlsSessionOptions &options = ClientTlsSessionOptions()) { using namespace tls; - if (!ctx) { return false; } + auto fail = [&](Error error, int ssl_error, uint64_t backend_error) { + if (out_error) { + out_error->error = error; + out_error->ssl_error = ssl_error; + out_error->backend_error = backend_error; + } + return false; + }; - bool is_ip = is_ip_address(host); + if (!ctx) { + session = nullptr; + return fail(Error::SSLConnection, 0, 0); + } #if defined(CPPHTTPLIB_MBEDTLS_SUPPORT) || defined(CPPHTTPLIB_WOLFSSL_SUPPORT) - // Chain verification happens during the handshake even for IP hosts; the - // certificate identity is verified post-handshake via verify_hostname() + // Mbed TLS and wolfSSL need the verification mode set explicitly; OpenSSL + // uses SSL_VERIFY_NONE and does all verification post-handshake. Chain + // verification happens during the handshake even for IP hosts; the + // certificate identity is verified post-handshake via verify_hostname(). set_verify_client(ctx, server_certificate_verification); #endif - session = create_session(ctx, sock); - if (!session) { return false; } + { + std::unique_lock guard; + if (options.ctx_mutex) { + guard = std::unique_lock(*options.ctx_mutex); + } + session = create_session(ctx, sock); + } + if (!session) { return fail(Error::SSLConnection, 0, get_error()); } - // RFC 6066: SNI must not be set for IP addresses. On Mbed TLS and wolfSSL - // set_hostname also sets SNI, so it must be skipped for IP hosts as well; - // their identity is checked post-handshake below instead. - if (!is_ip) { - if (server_certificate_verification) { - set_hostname(session, host.c_str()); - } else { - set_sni(session, host.c_str()); + // RFC 6066: SNI must not be set for IP addresses; skip it for IP hosts, so + // their identity is checked post-handshake below instead. On Mbed TLS and + // wolfSSL, set_sni also drives handshake-time hostname verification, so + // options.server_hostname_verification is threaded through here. + if (!is_ip_address(host)) { + if (!set_sni(session, host.c_str(), options.server_hostname_verification)) { + return fail(Error::SSLConnection, 0, get_error()); } } - if (!connect_nonblocking(session, sock, timeout_sec, timeout_usec, nullptr)) { - return false; + TlsError tls_err; + if (!connect_nonblocking(session, sock, timeout_sec, timeout_usec, + &tls_err)) { + auto error = Error::SSLConnection; + if (tls_err.code == ErrorCode::CertVerifyFailed) { + error = Error::SSLServerVerification; + } else if (tls_err.code == ErrorCode::HostnameMismatch) { + error = Error::SSLServerHostnameVerification; + } + return fail(error, static_cast(tls_err.code), tls_err.backend_code); } - if (server_certificate_verification) { - if (get_verify_result(session) != 0) { return false; } + auto verification_status = SSLVerifierResponse::NoDecisionMade; + if (options.session_verifier) { + verification_status = options.session_verifier(session); + } + + if (verification_status == SSLVerifierResponse::CertificateRejected) { + return fail(Error::SSLServerVerification, 0, get_error()); + } + + if (verification_status == SSLVerifierResponse::NoDecisionMade && + server_certificate_verification) { + auto verify_result = get_verify_result(session); + if (verify_result != 0) { + return fail(Error::SSLServerVerification, 0, + static_cast(verify_result)); + } + + auto server_cert = get_peer_cert(session); + if (!server_cert) { + return fail(Error::SSLServerVerification, 0, get_error()); + } + auto cert_guard = detail::scope_exit([&] { free_cert(server_cert); }); // Identity check against the peer certificate, post-handshake for all - // backends (same as SSLClient). For IP hosts this is the only identity - // verification since no hostname is bound during the handshake. - auto server_cert = get_peer_cert(session); - if (!server_cert) { return false; } - auto cert_guard = detail::scope_exit([&] { free_cert(server_cert); }); - if (!verify_hostname(server_cert, host.c_str())) { return false; } + // backends. For IP hosts this is the only identity verification, since no + // hostname is bound during the handshake. + if (options.server_hostname_verification) { + if (!verify_hostname(server_cert, host.c_str())) { + return fail(Error::SSLServerHostnameVerification, 0, + hostname_mismatch_code()); + } + } + +#ifdef CPPHTTPLIB_WINDOWS_AUTOMATIC_ROOT_CERTIFICATES_UPDATE + // Additional Windows Schannel verification. + // This provides real-time certificate validation with Windows Update + // integration, working with both OpenSSL and MbedTLS backends. + if (options.windows_cert_verification) { + std::vector der; + if (get_cert_der(server_cert, der)) { + uint64_t wincrypt_error = 0; + if (!verify_cert_with_windows_schannel( + der, host, options.server_hostname_verification, + wincrypt_error)) { + return fail(Error::SSLServerVerification, 0, wincrypt_error); + } + } + } +#endif } return true; @@ -5833,6 +6017,7 @@ std::string to_string(const Error error) { case Error::HTTPParsing: return "HTTP parsing failed"; case Error::InvalidRangeHeader: return "Invalid Range header"; case Error::UnsupportedContentEncoding: return "Unsupported Content-Encoding"; + case Error::WebSocketHandshake: return "WebSocket handshake failed"; default: break; } @@ -7039,6 +7224,26 @@ make_host_and_port_string_always_port(const std::string &host, int port) { return prepare_host_string(host) + ":" + std::to_string(port); } +// Value for the Host header a client sends when the caller supplied none. +// Only the value: callers decide where in their header list it goes. +std::string make_default_host_header_value(const std::string &host, + int port, bool is_ssl, + int address_family) { + if (address_family == AF_UNIX) { return "localhost"; } + return make_host_and_port_string(host, port, is_ssl); +} + +void add_default_user_agent_header(Request &req) { +#ifndef CPPHTTPLIB_NO_DEFAULT_USER_AGENT + if (!req.has_header("User-Agent")) { + req.set_header("User-Agent", + std::string("cpp-httplib/") + CPPHTTPLIB_VERSION); + } +#else + (void)req; +#endif +} + bool parse_no_proxy_entry(const std::string &token, NoProxyEntry &out); NormalizedTarget normalize_target(const std::string &host); bool ip_in_cidr(const IPBytes &ip, const IPBytes &net, int prefix_bits); @@ -8094,8 +8299,14 @@ bool Server::read_content_core( bool Server::handle_file_request(Request &req, Response &res) { for (const auto &entry : base_dirs_) { - // Prefix match - if (!req.path.compare(0, entry.mount_point.size(), entry.mount_point)) { + // Prefix match, on a path segment boundary. A mount point of "/mount" + // covers "/mount" and "/mount/...", but must not swallow "/mountdir/...". + // One that already ends in '/' (the root mount among them) carries its own + // boundary; set_mount_point() guarantees the mount point is not empty. + if (!req.path.compare(0, entry.mount_point.size(), entry.mount_point) && + (entry.mount_point.back() == '/' || + req.path.size() == entry.mount_point.size() || + req.path[entry.mount_point.size()] == '/')) { std::string sub_path = "/" + req.path.substr(entry.mount_point.size()); if (detail::is_valid_path(sub_path)) { auto path = entry.base_dir + sub_path; @@ -8985,8 +9196,7 @@ bool Server::process_and_close_socket(socket_t sock) { nullptr, &websocket_upgraded); }); - detail::shutdown_socket(sock); - detail::close_socket(sock); + detail::drain_and_close_socket(sock); return ret; } @@ -9198,29 +9408,20 @@ bool ClientImpl::read_response_line(Stream &strm, const Request &req, if (!line_reader.getline()) { return false; } -#ifdef CPPHTTPLIB_ALLOW_LF_AS_LINE_TERMINATOR - thread_local const std::regex re("(HTTP/1\\.[01]) (\\d{3})(?: (.*?))?\r?\n"); -#else - thread_local const std::regex re("(HTTP/1\\.[01]) (\\d{3})(?: (.*?))?\r\n"); -#endif - - std::cmatch m; - if (!std::regex_match(line_reader.ptr(), m, re)) { + if (!detail::parse_status_line(line_reader.ptr(), res.version, res.status, + res.reason)) { return req.method == "CONNECT"; } - res.version = std::string(m[1]); - res.status = std::stoi(std::string(m[2])); - res.reason = std::string(m[3]); // Ignore '100 Continue' (only when not using Expect: 100-continue explicitly) while (skip_100_continue && res.status == StatusCode::Continue_100) { if (!line_reader.getline()) { return false; } // CRLF if (!line_reader.getline()) { return false; } // next response line - if (!std::regex_match(line_reader.ptr(), m, re)) { return false; } - res.version = std::string(m[1]); - res.status = std::stoi(std::string(m[2])); - res.reason = std::string(m[3]); + if (!detail::parse_status_line(line_reader.ptr(), res.version, res.status, + res.reason)) { + return false; + } } return true; @@ -9356,12 +9557,9 @@ void ClientImpl::prepare_default_headers(Request &r, bool for_stream, // RFC 9110 5.3 recommends sending control data such as Host first, so // prepend it rather than appending it after the caller's own fields. if (!r.has_header("Host")) { - if (address_family_ == AF_UNIX) { - r.headers.emplace_front("Host", "localhost"); - } else { - r.headers.emplace_front( - "Host", detail::make_host_and_port_string(host_, port_, is_ssl())); - } + r.headers.emplace_front( + "Host", detail::make_default_host_header_value(host_, port_, is_ssl(), + address_family_)); } if (!r.has_header("Accept")) { r.headers.emplace("Accept", "*/*"); } @@ -9383,12 +9581,7 @@ void ClientImpl::prepare_default_headers(Request &r, bool for_stream, r.set_header("Accept-Encoding", accept_encoding); } -#ifndef CPPHTTPLIB_NO_DEFAULT_USER_AGENT - if (!r.has_header("User-Agent")) { - auto agent = std::string("cpp-httplib/") + CPPHTTPLIB_VERSION; - r.set_header("User-Agent", agent); - } -#endif + detail::add_default_user_agent_header(r); } if (!r.body.empty()) { @@ -12684,6 +12877,8 @@ void SSLClient::load_ca_cert_store(const char *ca_cert, bool SSLClient::load_certs() { auto ret = true; + // call_once rather than the plain flag WebSocketClient::create_stream() uses: + // one client is shared across concurrent requests here. std::call_once(initialize_cert_, [&]() { std::lock_guard guard(ctx_mutex_); @@ -12697,8 +12892,6 @@ bool SSLClient::load_certs() { } bool SSLClient::initialize_ssl(Socket &socket, Error &error) { - using namespace tls; - // Load CA certificates if server verification is enabled if (server_certificate_verification_) { if (!load_certs()) { @@ -12708,134 +12901,40 @@ bool SSLClient::initialize_ssl(Socket &socket, Error &error) { } } - bool is_ip = detail::is_ip_address(host_); - -#if defined(CPPHTTPLIB_MBEDTLS_SUPPORT) || defined(CPPHTTPLIB_WOLFSSL_SUPPORT) - // MbedTLS/wolfSSL need explicit verification mode (OpenSSL uses - // SSL_VERIFY_NONE by default and performs all verification post-handshake). - // Chain verification happens during the handshake even for IP hosts; the - // certificate identity is verified post-handshake via verify_hostname(). - set_verify_client(ctx_, server_certificate_verification_); + detail::ClientTlsSessionOptions options; + options.server_hostname_verification = server_hostname_verification_; + options.session_verifier = session_verifier_; + options.ctx_mutex = &ctx_mutex_; +#ifdef CPPHTTPLIB_WINDOWS_AUTOMATIC_ROOT_CERTIFICATES_UPDATE + // Skip Schannel when a custom CA cert is specified, as the Windows + // certificate store would not know about user-provided CA certificates. + // Also skip when system CA trust is explicitly disabled. + options.windows_cert_verification = + enable_windows_cert_verification_ && + system_ca_mode_ != SystemCAMode::Disabled && ca_cert_file_path_.empty() && + ca_cert_dir_path_.empty() && ca_cert_pem_.empty() && !ca_cert_store_set_; #endif - // Create TLS session - session_t session = nullptr; - { - std::lock_guard guard(ctx_mutex_); - session = create_session(ctx_, socket.sock); - } - - if (!session) { - error = Error::SSLConnection; - last_backend_error_ = get_error(); - return false; - } + tls::session_t session = nullptr; // Use scope_exit to ensure session is freed on error paths bool success = false; auto session_guard = detail::scope_exit([&] { - if (!success) { free_session(session); } + if (!success) { tls::free_session(session); } }); - // Set SNI extension (skip for IP addresses per RFC 6066). - // On MbedTLS, set_sni also enables hostname verification internally. - // On OpenSSL, set_sni only sets SNI; verification is done post-handshake. - if (!is_ip) { - if (!set_sni(session, host_.c_str())) { - error = Error::SSLConnection; - last_backend_error_ = get_error(); - return false; - } - } - - // Perform non-blocking TLS handshake with timeout - TlsError tls_err; - if (!connect_nonblocking(session, socket.sock, connection_timeout_sec_, - connection_timeout_usec_, &tls_err)) { - last_ssl_error_ = static_cast(tls_err.code); - last_backend_error_ = tls_err.backend_code; - if (tls_err.code == ErrorCode::CertVerifyFailed) { - error = Error::SSLServerVerification; - } else if (tls_err.code == ErrorCode::HostnameMismatch) { - error = Error::SSLServerHostnameVerification; - } else { - error = Error::SSLConnection; - } + detail::ClientTlsSessionError tls_error; + if (!detail::setup_client_tls_session( + host_, ctx_, session, socket.sock, server_certificate_verification_, + connection_timeout_sec_, connection_timeout_usec_, &tls_error, + options)) { + error = tls_error.error; + last_ssl_error_ = tls_error.ssl_error; + last_backend_error_ = tls_error.backend_error; output_error_log(error, nullptr); return false; } - // Post-handshake session verifier callback - auto verification_status = SSLVerifierResponse::NoDecisionMade; - if (session_verifier_) { verification_status = session_verifier_(session); } - - if (verification_status == SSLVerifierResponse::CertificateRejected) { - last_backend_error_ = get_error(); - error = Error::SSLServerVerification; - output_error_log(error, nullptr); - return false; - } - - // Default server certificate verification - if (verification_status == SSLVerifierResponse::NoDecisionMade && - server_certificate_verification_) { - verify_result_ = tls::get_verify_result(session); - if (verify_result_ != 0) { - last_backend_error_ = static_cast(verify_result_); - error = Error::SSLServerVerification; - output_error_log(error, nullptr); - return false; - } - - auto server_cert = get_peer_cert(session); - if (!server_cert) { - last_backend_error_ = get_error(); - error = Error::SSLServerVerification; - output_error_log(error, nullptr); - return false; - } - auto cert_guard = detail::scope_exit([&] { free_cert(server_cert); }); - - // Hostname verification (post-handshake for all cases). - // On OpenSSL, verification is always post-handshake (SSL_VERIFY_NONE). - // On MbedTLS, set_sni already enabled hostname verification during - // handshake for non-IP hosts, but this check is still needed for IP - // addresses where SNI is not set. - if (server_hostname_verification_) { - if (!verify_hostname(server_cert, host_.c_str())) { - last_backend_error_ = hostname_mismatch_code(); - error = Error::SSLServerHostnameVerification; - output_error_log(error, nullptr); - return false; - } - } - -#ifdef CPPHTTPLIB_WINDOWS_AUTOMATIC_ROOT_CERTIFICATES_UPDATE - // Additional Windows Schannel verification. - // This provides real-time certificate validation with Windows Update - // integration, working with both OpenSSL and MbedTLS backends. - // Skip when a custom CA cert is specified, as the Windows certificate - // store would not know about user-provided CA certificates. Also skip - // when system CA trust is explicitly disabled. - if (enable_windows_cert_verification_ && - system_ca_mode_ != SystemCAMode::Disabled && - ca_cert_file_path_.empty() && ca_cert_dir_path_.empty() && - ca_cert_pem_.empty() && !ca_cert_store_set_) { - std::vector der; - if (get_cert_der(server_cert, der)) { - uint64_t wincrypt_error = 0; - if (!detail::verify_cert_with_windows_schannel( - der, host_, server_hostname_verification_, wincrypt_error)) { - last_backend_error_ = wincrypt_error; - error = Error::SSLServerVerification; - output_error_log(error, nullptr); - return false; - } - } - } -#endif - } - success = true; socket.ssl = session; return true; @@ -13550,12 +13649,15 @@ void free_session(session_t session) { if (session) { SSL_free(static_cast(session)); } } -bool set_sni(session_t session, const char *hostname) { +bool set_sni(session_t session, const char *hostname, + bool /*verify_hostname*/) { if (!session || !hostname) return false; auto ssl = static_cast(session); - // Set SNI (Server Name Indication) only - does not enable verification + // Set SNI (Server Name Indication) only - does not enable verification. + // OpenSSL never binds identity checking to SNI (that happens post- + // handshake in setup_client_tls_session()), so verify_hostname is unused. #if defined(OPENSSL_IS_BORINGSSL) return SSL_set_tlsext_host_name(ssl, hostname) == 1; #else @@ -13565,32 +13667,6 @@ bool set_sni(session_t session, const char *hostname) { #endif } -bool set_hostname(session_t session, const char *hostname) { - if (!session || !hostname) return false; - - auto ssl = static_cast(session); - - // Enable hostname verification - auto param = SSL_get0_param(ssl); - if (!param) return false; - - if (detail::is_ip_address(hostname)) { - // RFC 6066: SNI must not be set for IP addresses; verify against the - // certificate's IP SANs instead of its DNS names - if (X509_VERIFY_PARAM_set1_ip_asc(param, hostname) != 1) { return false; } - } else { - // Set SNI (Server Name Indication) - if (!set_sni(session, hostname)) { return false; } - - X509_VERIFY_PARAM_set_hostflags(param, - X509_CHECK_FLAG_NO_PARTIAL_WILDCARDS); - if (X509_VERIFY_PARAM_set1_host(param, hostname, 0) != 1) { return false; } - } - - SSL_set_verify(ssl, SSL_VERIFY_PEER, nullptr); - return true; -} - TlsError connect(session_t session) { if (!session) { return TlsError(); } @@ -14219,6 +14295,21 @@ struct MbedTlsSession { unsigned char peeked_byte = 0; bool has_peeked_byte = false; + // Set by set_sni() when the caller disabled hostname verification, so the + // verify callback can clear the CN/SAN mismatch flag while still enforcing + // the rest of the chain (Mbed TLS ties SNI and identity checking together; + // OpenSSL and wolfSSL keep them independent). + bool suppress_hostname_mismatch = false; + + // Copied from the owning MbedTlsContext at creation. set_sni() uses this to + // decide which verify callback to install when hostname verification is + // disabled: mbedtls_verify_callback() when a user callback is genuinely + // wired for this context, or a self-contained one otherwise, so a session + // that never opted into a callback never consults the process-wide + // set_verify_callback() slot (which some other, unrelated client may have + // populated). + bool has_verify_callback = false; + MbedTlsSession() { mbedtls_ssl_init(&ssl); } ~MbedTlsSession() { mbedtls_ssl_free(&ssl); } @@ -14235,7 +14326,8 @@ int &mbedtls_last_error() { } // Helper to map Mbed TLS error to ErrorCode -ErrorCode map_mbedtls_error(int ret, int &out_errno) { +ErrorCode map_mbedtls_error(int ret, int &out_errno, + uint32_t verify_flags) { if (ret == 0) { return ErrorCode::Success; } if (ret == MBEDTLS_ERR_SSL_WANT_READ) { return ErrorCode::WantRead; } if (ret == MBEDTLS_ERR_SSL_WANT_WRITE) { return ErrorCode::WantWrite; } @@ -14248,11 +14340,34 @@ ErrorCode map_mbedtls_error(int ret, int &out_errno) { return ErrorCode::SyscallError; } if (ret == MBEDTLS_ERR_X509_CERT_VERIFY_FAILED) { + // Unlike OpenSSL/wolfSSL, Mbed TLS folds the CN/SAN identity check into + // the handshake's chain verification (see set_sni()); a mismatch there + // is reported the same way as any other verify_flags bit. Report it as + // HostnameMismatch, matching the other backends and the post-handshake + // identity check below, but only when naming is the sole problem - + // if the chain itself is also untrusted/expired/etc., that takes + // priority over the naming detail. + if (verify_flags == static_cast(hostname_mismatch_code())) { + return ErrorCode::HostnameMismatch; + } return ErrorCode::CertVerifyFailed; } return ErrorCode::Fatal; } +// Populates a TlsError from a failed (non-zero) mbedtls_ssl_handshake() +// return value, including the verify-flags-dependent HostnameMismatch +// mapping; shared by connect() and connect_nonblocking() so the +// backend_code policy for that mapping only lives in one place. +void fill_mbedtls_tls_error(TlsError &err, mbedtls_ssl_context &ssl, + int ret) { + auto verify_flags = mbedtls_ssl_get_verify_result(&ssl); + err.code = map_mbedtls_error(ret, err.sys_errno, verify_flags); + err.backend_code = err.code == ErrorCode::HostnameMismatch + ? static_cast(verify_flags) + : static_cast(-ret); +} + // A TLS 1.3 NewSessionTicket (signaled by default on Mbed TLS 4.x) is a // non-fatal notification delivered between records, not an error and not // application data, so I/O calls that see it should just be retried. Kept in @@ -14362,18 +14477,44 @@ int mbedtls_sni_callback(void *p_ctx, mbedtls_ssl_context *ssl, return 0; // Accept any SNI } +void mbedtls_clear_cn_mismatch(uint32_t *flags) { + *flags &= ~static_cast(hostname_mismatch_code()); +} + +// Verify callback used when hostname verification is disabled for a session +// that has no user-supplied verify callback of its own (MbedTlsSession:: +// has_verify_callback is false). Deliberately does not consult +// get_verify_callback(): that slot is process-wide, so reading it here would +// pick up whatever another, unrelated client last installed there. +int mbedtls_mask_hostname_mismatch_callback(void *data, + mbedtls_x509_crt *, int, + uint32_t *flags) { + (void)data; + mbedtls_clear_cn_mismatch(flags); + return 0; +} + int mbedtls_verify_callback(void *data, mbedtls_x509_crt *crt, int cert_depth, uint32_t *flags); // MbedTLS verify callback wrapper int mbedtls_verify_callback(void *data, mbedtls_x509_crt *crt, int cert_depth, uint32_t *flags) { - auto &callback = get_verify_callback(); - if (!callback) { return 0; } // Continue with default verification - // data points to the MbedTlsSession auto *session = static_cast(data); + // set_sni() disabled hostname verification for this session: drop the + // CN/SAN mismatch flag so it doesn't fail the chain check below, mirroring + // the OpenSSL/wolfSSL backends where identity checking is independent of + // SNI. The final pass/fail decision still comes from the remaining flags + // (or, below, from the user's own verify callback). + if (session && session->suppress_hostname_mismatch) { + mbedtls_clear_cn_mismatch(flags); + } + + auto &callback = get_verify_callback(); + if (!callback) { return 0; } // Continue with default verification + // Build context VerifyContext verify_ctx; verify_ctx.session = static_cast(session); @@ -14789,6 +14930,7 @@ session_t create_session(ctx_t ctx, socket_t sock) { // Set per-session verify callback with session pointer if callback is // registered + session->has_verify_callback = mctx->has_verify_callback; if (mctx->has_verify_callback) { mbedtls_ssl_set_verify(&session->ssl, impl::mbedtls_verify_callback, session); @@ -14801,10 +14943,15 @@ void free_session(session_t session) { if (session) { delete static_cast(session); } } -bool set_sni(session_t session, const char *hostname) { +bool set_sni(session_t session, const char *hostname, + bool verify_hostname) { if (!session || !hostname) { return false; } auto msession = static_cast(session); + // mbedtls_ssl_set_hostname() both sends the SNI extension and binds the + // handshake-time CN/SAN check to `hostname`; the two can't be requested + // independently, so a disabled hostname check is handled below by masking + // the resulting mismatch flag instead of skipping this call. int ret = mbedtls_ssl_set_hostname(&msession->ssl, hostname); if (ret != 0) { impl::mbedtls_last_error() = ret; @@ -14812,12 +14959,22 @@ bool set_sni(session_t session, const char *hostname) { } msession->hostname = hostname; - return true; -} -bool set_hostname(session_t session, const char *hostname) { - // In Mbed TLS, set_hostname also sets up hostname verification - return set_sni(session, hostname); + if (!verify_hostname) { + msession->suppress_hostname_mismatch = true; + // If a user verify callback is already wired for this session, + // mbedtls_verify_callback() masks the mismatch flag itself before + // consulting it (see suppress_hostname_mismatch above) - reinstalling it + // here would be redundant. Otherwise install the self-contained masking + // callback, which never touches the process-wide callback slot. + if (!msession->has_verify_callback) { + mbedtls_ssl_set_verify(&msession->ssl, + impl::mbedtls_mask_hostname_mismatch_callback, + msession); + } + } + + return true; } TlsError connect(session_t session) { @@ -14836,8 +14993,7 @@ TlsError connect(session_t session) { if (ret == 0) { err.code = ErrorCode::Success; } else { - err.code = impl::map_mbedtls_error(ret, err.sys_errno); - err.backend_code = static_cast(-ret); + impl::fill_mbedtls_tls_error(err, msession->ssl, ret); impl::mbedtls_last_error() = ret; } @@ -14888,10 +15044,7 @@ bool connect_nonblocking(session_t session, socket_t sock, } // TlsError or timeout - if (err) { - err->code = impl::map_mbedtls_error(ret, err->sys_errno); - err->backend_code = static_cast(-ret); - } + if (err) { impl::fill_mbedtls_tls_error(*err, msession->ssl, ret); } impl::mbedtls_last_error() = ret; return false; } @@ -14957,7 +15110,7 @@ ssize_t read(session_t session, void *buf, size_t len, TlsError &err) { return 0; } - err.code = impl::map_mbedtls_error(ret, err.sys_errno); + err.code = impl::map_mbedtls_error(ret, err.sys_errno, 0); err.backend_code = static_cast(-ret); impl::mbedtls_last_error() = ret; // mbedTLS signals a clean close_notify via a negative error code rather @@ -14990,7 +15143,7 @@ ssize_t write(session_t session, const void *buf, size_t len, return 0; } - err.code = impl::map_mbedtls_error(ret, err.sys_errno); + err.code = impl::map_mbedtls_error(ret, err.sys_errno, 0); err.backend_code = static_cast(-ret); impl::mbedtls_last_error() = ret; return -1; @@ -15953,7 +16106,8 @@ void free_session(session_t session) { if (session) { delete static_cast(session); } } -bool set_sni(session_t session, const char *hostname) { +bool set_sni(session_t session, const char *hostname, + bool verify_hostname) { if (!session || !hostname) { return false; } auto wsession = static_cast(session); @@ -15965,18 +16119,15 @@ bool set_sni(session_t session, const char *hostname) { return false; } - // Also set hostname for verification - wolfSSL_check_domain_name(wsession->ssl, hostname); + // wolfSSL_check_domain_name binds identity checking to the handshake, + // separately from the SNI extension sent above; skip it when hostname + // verification is disabled so only the chain is checked, matching OpenSSL. + if (verify_hostname) { wolfSSL_check_domain_name(wsession->ssl, hostname); } wsession->hostname = hostname; return true; } -bool set_hostname(session_t session, const char *hostname) { - // In wolfSSL, set_hostname also sets up hostname verification - return set_sni(session, hostname); -} - TlsError connect(session_t session) { TlsError err; if (!session) { @@ -16922,6 +17073,24 @@ WebSocketClient::WebSocketClient( } } +#ifdef CPPHTTPLIB_SSL_ENABLED +WebSocketClient::WebSocketClient( + const std::string &scheme_host_port_path, const PemMemory &pem, + const Headers &headers) + : WebSocketClient(scheme_host_port_path, headers) { + // For ws:// URLs the client certificate is silently ignored, consistent + // with the TLS-only setters such as set_ca_cert_path(). + if (is_valid_ && is_ssl_ && pem.cert_pem && pem.key_pem) { + if (!tls::set_client_cert_pem(tls_ctx_, pem.cert_pem, pem.key_pem, + pem.private_key_password)) { + tls::free_context(tls_ctx_); + tls_ctx_ = nullptr; + is_valid_ = false; + } + } +} +#endif + WebSocketClient::~WebSocketClient() { shutdown_and_close(); #ifdef CPPHTTPLIB_SSL_ENABLED @@ -16956,21 +17125,33 @@ void WebSocketClient::shutdown_and_close() { } } -bool WebSocketClient::create_stream(std::unique_ptr &strm) { +bool WebSocketClient::create_stream(std::unique_ptr &strm, + Error &error, int &ssl_error, + uint64_t &ssl_backend_error) { #ifdef CPPHTTPLIB_SSL_ENABLED if (is_ssl_) { + // A plain flag rather than SSLClient::load_certs()'s call_once: connect() + // is not safe to call concurrently on one client to begin with, since + // nothing else here is guarded either. if (server_certificate_verification_ && !certs_loaded_) { uint64_t backend_error = 0; - detail::load_client_ca_config(tls_ctx_, ca_cert_file_path_, std::string(), - custom_ca_loaded_, system_ca_mode_, - backend_error); + detail::load_client_ca_config(tls_ctx_, ca_cert_file_path_, + ca_cert_dir_path_, custom_ca_loaded_, + system_ca_mode_, backend_error); certs_loaded_ = true; } + detail::ClientTlsSessionOptions options; + options.server_hostname_verification = server_hostname_verification_; + + detail::ClientTlsSessionError tls_error; if (!detail::setup_client_tls_session(host_, tls_ctx_, tls_session_, sock_, server_certificate_verification_, - read_timeout_sec_, - read_timeout_usec_)) { + read_timeout_sec_, read_timeout_usec_, + &tls_error, options)) { + error = tls_error.error; + ssl_error = tls_error.ssl_error; + ssl_backend_error = tls_error.backend_error; return false; } @@ -16979,6 +17160,10 @@ bool WebSocketClient::create_stream(std::unique_ptr &strm) { write_timeout_sec_, write_timeout_usec_)); return true; } +#else + (void)error; + (void)ssl_error; + (void)ssl_backend_error; #endif strm = std::unique_ptr( new detail::SocketStream(sock_, read_timeout_sec_, read_timeout_usec_, @@ -16994,24 +17179,15 @@ void WebSocketClient::prepare_default_headers(Request &req) { #endif if (!req.has_header("Host")) { - if (address_family_ == AF_UNIX) { - req.headers.emplace("Host", "localhost"); - } else { - req.headers.emplace( - "Host", detail::make_host_and_port_string(host_, port_, is_ssl)); - } + req.headers.emplace("Host", detail::make_default_host_header_value( + host_, port_, is_ssl, address_family_)); } -#ifndef CPPHTTPLIB_NO_DEFAULT_USER_AGENT - if (!req.has_header("User-Agent")) { - auto agent = std::string("cpp-httplib/") + CPPHTTPLIB_VERSION; - req.set_header("User-Agent", agent); - } -#endif + detail::add_default_user_agent_header(req); } -bool WebSocketClient::connect() { - if (!is_valid_) { return false; } +Result WebSocketClient::connect() { + if (!is_valid_) { return Result{Error::Connection, -1, Headers{}}; } shutdown_and_close(); // Check is custom IP or hostname specified for host_ @@ -17019,19 +17195,29 @@ bool WebSocketClient::connect() { std::string ip; detail::apply_addr_map(addr_map_, host_, connect_host, ip); - Error error; + auto error = Error::Success; sock_ = detail::create_client_socket( connect_host, ip, port_, address_family_, tcp_nodelay_, ipv6_v6only_, socket_options_, connection_timeout_sec_, connection_timeout_usec_, read_timeout_sec_, read_timeout_usec_, write_timeout_sec_, write_timeout_usec_, interface_, error); - if (sock_ == INVALID_SOCKET) { return false; } + if (sock_ == INVALID_SOCKET) { + if (error == Error::Success) { error = Error::Connection; } + return Result{error, -1, Headers{}}; + } std::unique_ptr strm; - if (!create_stream(strm)) { + auto stream_error = Error::SSLConnection; + int ssl_error = 0; + uint64_t ssl_backend_error = 0; + if (!create_stream(strm, stream_error, ssl_error, ssl_backend_error)) { shutdown_and_close(); - return false; +#ifdef CPPHTTPLIB_SSL_ENABLED + return Result{stream_error, -1, Headers{}, ssl_error, ssl_backend_error}; +#else + return Result{stream_error, -1, Headers{}}; +#endif } Request req; @@ -17040,17 +17226,17 @@ bool WebSocketClient::connect() { req.headers = headers_; prepare_default_headers(req); - std::string selected_subprotocol; - if (!detail::perform_websocket_handshake(*strm, req, selected_subprotocol)) { + detail::WebSocketUpgradeResponse upgrade; + if (!detail::perform_websocket_handshake(*strm, req, upgrade)) { shutdown_and_close(); - return false; + return Result{upgrade.error, upgrade.status, std::move(upgrade.headers)}; } - subprotocol_ = std::move(selected_subprotocol); + subprotocol_ = std::move(upgrade.selected_subprotocol); ws_ = std::unique_ptr(new WebSocket(std::move(strm), req, false, websocket_ping_interval_sec_, websocket_max_missed_pongs_)); - return true; + return Result{Error::Success, upgrade.status, std::move(upgrade.headers)}; } ReadResult WebSocketClient::read(std::string &msg) { @@ -17125,8 +17311,11 @@ void WebSocketClient::set_hostname_addr_map( #ifdef CPPHTTPLIB_SSL_ENABLED -void WebSocketClient::set_ca_cert_path(const std::string &path) { - ca_cert_file_path_ = path; +void +WebSocketClient::set_ca_cert_path(const std::string &ca_cert_file_path, + const std::string &ca_cert_dir_path) { + ca_cert_file_path_ = ca_cert_file_path; + ca_cert_dir_path_ = ca_cert_dir_path; } void WebSocketClient::set_ca_cert_store(tls::ca_store_t store) { @@ -17152,6 +17341,10 @@ WebSocketClient::enable_server_certificate_verification(bool enabled) { server_certificate_verification_ = enabled; } +void WebSocketClient::enable_server_hostname_verification(bool enabled) { + server_hostname_verification_ = enabled; +} + void WebSocketClient::enable_system_ca(bool enabled) { system_ca_mode_ = enabled ? SystemCAMode::Enabled : SystemCAMode::Disabled; } diff --git a/vendor/cpp-httplib/httplib.h b/vendor/cpp-httplib/httplib.h index 94605defa..f82e58e49 100644 --- a/vendor/cpp-httplib/httplib.h +++ b/vendor/cpp-httplib/httplib.h @@ -8,8 +8,8 @@ #ifndef CPPHTTPLIB_HTTPLIB_H #define CPPHTTPLIB_HTTPLIB_H -#define CPPHTTPLIB_VERSION "0.52.0" -#define CPPHTTPLIB_VERSION_NUM "0x003400" +#define CPPHTTPLIB_VERSION "0.53.0" +#define CPPHTTPLIB_VERSION_NUM "0x003500" #ifdef _WIN32 #if defined(_WIN32_WINNT) && _WIN32_WINNT < 0x0A00 @@ -1805,6 +1805,7 @@ enum class Error { HTTPParsing, InvalidRangeHeader, UnsupportedContentEncoding, + WebSocketHandshake, // For internal use only SSLPeerCouldBeClosed_, @@ -3233,8 +3234,6 @@ private: // Used to keep custom CA configuration exclusive with system CA loading. bool ca_cert_store_set_ = false; - long verify_result_ = 0; - std::function session_verifier_; #ifdef CPPHTTPLIB_WINDOWS_AUTOMATIC_ROOT_CERTIFICATES_UPDATE @@ -4204,6 +4203,50 @@ enum class CloseStatus : uint16_t { enum ReadResult : int { Fail = 0, Text = 1, Binary = 2 }; +// Result of WebSocketClient::connect(). Truthy only when the WebSocket +// upgrade handshake fully succeeded. On failure error() identifies the +// failing layer; status()/headers() expose the server's upgrade response +// when one was received (status() is -1 otherwise). +class Result { +public: + Result() = default; + Result(Error err, int status, Headers &&headers) + : err_(err), status_(status), headers_(std::move(headers)) {} + + explicit operator bool() const { return err_ == Error::Success; } + Error error() const { return err_; } + + // Upgrade response info + int status() const { return status_; } + const Headers &headers() const { return headers_; } + std::string get_header_value(const std::string &key, + const char *def = "") const { + return detail::get_header_value(headers_, key, def, 0); + } + bool has_header(const std::string &key) const { + return headers_.find(key) != headers_.end(); + } + +#ifdef CPPHTTPLIB_SSL_ENABLED + Result(Error err, int status, Headers &&headers, int ssl_error, + uint64_t ssl_backend_error) + : err_(err), status_(status), headers_(std::move(headers)), + ssl_error_(ssl_error), ssl_backend_error_(ssl_backend_error) {} + + int ssl_error() const { return ssl_error_; } + uint64_t ssl_backend_error() const { return ssl_backend_error_; } +#endif + +private: + Error err_ = Error::Unknown; // a default-constructed Result is falsy + int status_ = -1; + Headers headers_; +#ifdef CPPHTTPLIB_SSL_ENABLED + int ssl_error_ = 0; + uint64_t ssl_backend_error_ = 0; +#endif +}; + class WebSocket { public: WebSocket(const WebSocket &) = delete; @@ -4270,7 +4313,7 @@ public: bool is_valid() const; - bool connect(); + Result connect(); ReadResult read(std::string &msg); bool send(const std::string &data); bool send(const char *data, size_t len); @@ -4279,28 +4322,52 @@ public: bool is_open() const; const std::string &subprotocol() const; void set_read_timeout(time_t sec, time_t usec = 0); + template + void set_read_timeout(const std::chrono::duration &duration); + void set_write_timeout(time_t sec, time_t usec = 0); + template + void set_write_timeout(const std::chrono::duration &duration); + void set_websocket_ping_interval(time_t sec); void set_websocket_max_missed_pongs(int count); void set_tcp_nodelay(bool on); void set_address_family(int family); void set_ipv6_v6only(bool on); void set_socket_options(SocketOptions socket_options); + void set_connection_timeout(time_t sec, time_t usec = 0); + template + void + set_connection_timeout(const std::chrono::duration &duration); + void set_interface(const std::string &intf); void set_hostname_addr_map(std::map addr_map); #ifdef CPPHTTPLIB_SSL_ENABLED - void set_ca_cert_path(const std::string &path); + struct PemMemory { + const char *cert_pem; + size_t cert_pem_len; + const char *key_pem; + size_t key_pem_len; + const char *private_key_password; + }; + explicit WebSocketClient(const std::string &scheme_host_port_path, + const PemMemory &pem, const Headers &headers = {}); + + void set_ca_cert_path(const std::string &ca_cert_file_path, + const std::string &ca_cert_dir_path = std::string()); void set_ca_cert_store(tls::ca_store_t store); void load_ca_cert_store(const char *ca_cert, std::size_t size); void enable_server_certificate_verification(bool enabled); + void enable_server_hostname_verification(bool enabled); void enable_system_ca(bool enabled); #endif private: void shutdown_and_close(); - bool create_stream(std::unique_ptr &strm); + bool create_stream(std::unique_ptr &strm, Error &error, + int &ssl_error, uint64_t &ssl_backend_error); void prepare_default_headers(Request &req); std::string host_; @@ -4335,13 +4402,37 @@ private: tls::ctx_t tls_ctx_ = nullptr; tls::session_t tls_session_ = nullptr; std::string ca_cert_file_path_; + std::string ca_cert_dir_path_; bool custom_ca_loaded_ = false; bool certs_loaded_ = false; SystemCAMode system_ca_mode_ = SystemCAMode::Auto; bool server_certificate_verification_ = true; + bool server_hostname_verification_ = true; #endif }; +template +inline void WebSocketClient::set_read_timeout( + const std::chrono::duration &duration) { + detail::duration_to_sec_and_usec( + duration, [&](time_t sec, time_t usec) { set_read_timeout(sec, usec); }); +} + +template +inline void WebSocketClient::set_write_timeout( + const std::chrono::duration &duration) { + detail::duration_to_sec_and_usec( + duration, [&](time_t sec, time_t usec) { set_write_timeout(sec, usec); }); +} + +template +inline void WebSocketClient::set_connection_timeout( + const std::chrono::duration &duration) { + detail::duration_to_sec_and_usec(duration, [&](time_t sec, time_t usec) { + set_connection_timeout(sec, usec); + }); +} + namespace impl { bool is_valid_utf8(const std::string &s); diff --git a/vendor/sheredom/patch-bsd.patch b/vendor/sheredom/patch-bsd.patch deleted file mode 100644 index 2532050e2..000000000 --- a/vendor/sheredom/patch-bsd.patch +++ /dev/null @@ -1,19 +0,0 @@ -Fix build on NetBSD, which provides posix_spawn_file_actions_addchdir() -but not the _np() variant. - -Upstream PR: https://github.com/sheredom/subprocess.h/pull/102 -Applied locally by scripts/sync_vendor.py until it is merged upstream. - -diff --git a/subprocess.h b/subprocess.h -index 5e809023a4..74a4e006c7 100644 ---- a/subprocess.h -+++ b/subprocess.h -@@ -1205,7 +1205,7 @@ cleanup: - - // Set working directory - if (process_cwd) { --#if defined(__APPLE__) && MAC_OS_X_VERSION_MIN_REQUIRED >= 260000 -+#if defined(__NetBSD__) || (defined(__APPLE__) && MAC_OS_X_VERSION_MIN_REQUIRED >= 260000) - posix_error = posix_spawn_file_actions_addchdir(&actions, process_cwd); - #else - #if defined(__APPLE__) && defined(__clang__) diff --git a/vendor/sheredom/patch-glibc-older-than-2.29.patch b/vendor/sheredom/patch-glibc-older-than-2.29.patch deleted file mode 100644 index 914421309..000000000 --- a/vendor/sheredom/patch-glibc-older-than-2.29.patch +++ /dev/null @@ -1,47 +0,0 @@ -Fix building against glibc older than 2.29, which has no -posix_spawn_file_actions_addchdir_np (the symbol is genuinely absent from -libc.so, so no feature-test macro helps). Affects manylinux2014 (glibc 2.17) -and manylinux_2_28, and was reported on RHEL 8.1. A requested process_cwd now -fails with ENOSYS there instead of failing the build. - -Upstream PR: https://github.com/sheredom/subprocess.h/pull/104 -Applied locally by scripts/sync_vendor.py until it is merged upstream. - -(the README.md and test/ changes from the PR are omitted, we only vendor -subprocess.h; rebased on top of patch-bsd.patch, so apply it after that one) - -diff --git a/subprocess.h b/subprocess.h -index 1ef424a..c363393 100644 ---- a/subprocess.h -+++ b/subprocess.h -@@ -274,6 +274,21 @@ subprocess_weak int subprocess_alive(struct subprocess_s *const process); - #include - #endif - -+/* Whether subprocess_create_ex can honour process_cwd. glibc only gained -+ posix_spawn_file_actions_addchdir_np in 2.29. Define this yourself to -+ override the detection, for instance on musl older than 1.1.24. */ -+#if !defined(SUBPROCESS_HAVE_CWD) -+#if defined(__GLIBC__) -+#if __GLIBC_PREREQ(2, 29) -+#define SUBPROCESS_HAVE_CWD 1 -+#else -+#define SUBPROCESS_HAVE_CWD 0 -+#endif -+#else -+#define SUBPROCESS_HAVE_CWD 1 -+#endif -+#endif -+ - #if defined(_WIN32) - - #include -@@ -1219,6 +1234,8 @@ cleanup: - if (process_cwd) { - #if defined(__NetBSD__) || (defined(__APPLE__) && MAC_OS_X_VERSION_MIN_REQUIRED >= 260000) - posix_error = posix_spawn_file_actions_addchdir(&actions, process_cwd); -+#elif !SUBPROCESS_HAVE_CWD -+ posix_error = ENOSYS; - #else - #if defined(__APPLE__) && defined(__clang__) - #pragma clang diagnostic push diff --git a/vendor/sheredom/patch-windows-quote-backslash.patch b/vendor/sheredom/patch-windows-quote-backslash.patch deleted file mode 100644 index 0204746d9..000000000 --- a/vendor/sheredom/patch-windows-quote-backslash.patch +++ /dev/null @@ -1,107 +0,0 @@ -Fix Windows command line quoting of backslash runs: a trailing backslash, or -backslashes preceding a double quote, were not doubled, so CommandLineToArgvW -in the child parsed them as escapes and mangled the argument list. - -Upstream PR: https://github.com/sheredom/subprocess.h/pull/101 -Applied locally by scripts/sync_vendor.py until it is merged upstream. - -(the test/ changes from the PR are omitted, we only vendor subprocess.h) - -diff --git a/subprocess.h b/subprocess.h -index 5e80902..b06ad4d 100644 ---- a/subprocess.h -+++ b/subprocess.h -@@ -653,6 +653,7 @@ int subprocess_create_ex(const char *const commandLine[], int options, - int wide_len; - int i, j; - int need_quoting; -+ subprocess_size_t bs_run; - unsigned long flags = 0; - unsigned long last_error = 0; - int result = subprocess_error_unknown; -@@ -906,25 +907,29 @@ int subprocess_create_ex(const char *const commandLine[], int options, - len++; - - // Quote the argument if it has a space in it -- if (strpbrk(commandLine[i], "\t\v ") != SUBPROCESS_NULL || -- commandLine[i][0] == SUBPROCESS_NULL) -+ need_quoting = strpbrk(commandLine[i], "\t\v ") != SUBPROCESS_NULL || -+ commandLine[i][0] == SUBPROCESS_NULL; -+ if (need_quoting) - len += 2; - -+ bs_run = 0; - for (j = 0; '\0' != commandLine[i][j]; j++) { -- switch (commandLine[i][j]) { -- default: -- break; -- case '\\': -- if (commandLine[i][j + 1] == '"') { -- len++; -- } -+ len++; - -- break; -- case '"': -- len++; -- break; -+ if ('\\' == commandLine[i][j]) { -+ bs_run++; -+ } else { -+ if ('"' == commandLine[i][j]) { -+ // Duplicate the preceding run and escape the quote. -+ len += bs_run + 1; -+ } -+ bs_run = 0; - } -- len++; -+ } -+ -+ if (need_quoting) { -+ // Duplicate trailing slashes before the generated closing quote. -+ len += bs_run; - } - } - -@@ -949,22 +954,29 @@ int subprocess_create_ex(const char *const commandLine[], int options, - commandLineCombined[len++] = '"'; - } - -- for (j = 0; '\0' != commandLine[i][j]; j++) { -- switch (commandLine[i][j]) { -- default: -- break; -- case '\\': -- if (commandLine[i][j + 1] == '"') { -- commandLineCombined[len++] = '\\'; -- } -+ for (j = 0; '\0' != commandLine[i][j];) { -+ bs_run = 0; -+ while ('\\' == commandLine[i][j]) { -+ bs_run++; -+ j++; -+ } -+ -+ if ('"' == commandLine[i][j]) { -+ // 2n + 1 slashes preserve n slashes and escape the quote. -+ bs_run = (bs_run * 2) + 1; -+ } else if ('\0' == commandLine[i][j] && need_quoting) { -+ // 2n slashes preserve n slashes before the closing quote. -+ bs_run *= 2; -+ } - -- break; -- case '"': -+ while (bs_run > 0) { - commandLineCombined[len++] = '\\'; -- break; -+ bs_run--; - } - -- commandLineCombined[len++] = commandLine[i][j]; -+ if ('\0' != commandLine[i][j]) { -+ commandLineCombined[len++] = commandLine[i][j++]; -+ } - } - if (need_quoting) { - commandLineCombined[len++] = '"'; diff --git a/vendor/sheredom/subprocess.h b/vendor/sheredom/subprocess.h index c36339387..c3af8a498 100644 --- a/vendor/sheredom/subprocess.h +++ b/vendor/sheredom/subprocess.h @@ -107,7 +107,8 @@ enum subprocess_error_e { subprocess_error_permission_denied = -5, subprocess_error_no_memory = -6, subprocess_error_pipe = -7, - subprocess_error_spawn = -8 + subprocess_error_spawn = -8, + subprocess_error_not_supported = -9 }; #if defined(__cplusplus) @@ -275,8 +276,10 @@ subprocess_weak int subprocess_alive(struct subprocess_s *const process); #endif /* Whether subprocess_create_ex can honour process_cwd. glibc only gained - posix_spawn_file_actions_addchdir_np in 2.29. Define this yourself to - override the detection, for instance on musl older than 1.1.24. */ + posix_spawn_file_actions_addchdir_np in 2.29, and macOS in 10.15; the SDKs + mark it unavailable on iOS, tvOS and watchOS, where the undefined version + macro folds to 0 and so answers correctly. Define this yourself to override + the detection, for instance on musl older than 1.1.24. */ #if !defined(SUBPROCESS_HAVE_CWD) #if defined(__GLIBC__) #if __GLIBC_PREREQ(2, 29) @@ -284,11 +287,27 @@ subprocess_weak int subprocess_alive(struct subprocess_s *const process); #else #define SUBPROCESS_HAVE_CWD 0 #endif +#elif defined(__APPLE__) && MAC_OS_X_VERSION_MIN_REQUIRED < 101500 +#define SUBPROCESS_HAVE_CWD 0 #else #define SUBPROCESS_HAVE_CWD 1 #endif #endif +/* Whether posix_spawn reports a failed exec back to the caller. glibc only + started doing so in 2.24; before that the child silently exits with 127. */ +#if !defined(SUBPROCESS_SPAWN_REPORTS_EXEC_ERRORS) +#if defined(__GLIBC__) +#if __GLIBC_PREREQ(2, 24) +#define SUBPROCESS_SPAWN_REPORTS_EXEC_ERRORS 1 +#else +#define SUBPROCESS_SPAWN_REPORTS_EXEC_ERRORS 0 +#endif +#else +#define SUBPROCESS_SPAWN_REPORTS_EXEC_ERRORS 1 +#endif +#endif + #if defined(_WIN32) #include @@ -554,6 +573,8 @@ int subprocess_error_from_errno(int error) { case ENFILE: case ENOMEM: return subprocess_error_no_memory; + case ENOSYS: + return subprocess_error_not_supported; default: return subprocess_error_unknown; } @@ -1358,6 +1379,17 @@ cleanup: goto cleanup; } } else { +#if !SUBPROCESS_SPAWN_REPORTS_EXEC_ERRORS + /* posix_spawn cannot tell us the exec failed, so check up front */ + if (0 != access(commandLine[0], X_OK)) { + saved_errno = errno; + result = subprocess_error_from_errno(saved_errno); + if (subprocess_error_unknown == result) { + result = subprocess_error_spawn; + } + goto cleanup; + } +#endif posix_error = posix_spawn(&child, commandLine[0], &actions, SUBPROCESS_NULL, SUBPROCESS_CONST_CAST(char *const *, commandLine),