server: add --no-chat-input-special

common/arg.cpp

@@ -3828,6 +3828,15 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.prefill_assistant = value;
         }
     ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_PREFILL_ASSISTANT"));
+    add_opt(common_arg(
+        {"--chat-input-special"},
+        {"--no-chat-input-special"},
+        string_format("experimental: whether to parse special tokens from user input messages in chat completions, does not affect the raw /completions endpoint (default: %s)\n"
+                      "note: this is not a security feature, do not rely on it to prevent prompt injection", params.chat_input_special ? "enabled" : "disabled"),
+        [](common_params & params, bool value) {
+            params.chat_input_special = value;
+        }
+    ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_CHAT_INPUT_SPECIAL"));
     add_opt(common_arg(
         {"-sps", "--slot-prompt-similarity"}, "SIMILARITY",
         string_format("how much the prompt of a request must match the prompt of a slot in order to use that slot (default: %.2f, 0.0 = disabled)\n", params.slot_prompt_similarity),

common/common.h

@@ -679,6 +679,7 @@ struct common_params {
     common_reasoning_format reasoning_format = COMMON_REASONING_FORMAT_DEEPSEEK;
     int enable_reasoning = -1; // -1 = auto, 0 = disable, 1 = enable
     bool prefill_assistant = true; // if true, any trailing assistant message will be prefilled into the response
+    bool chat_input_special = true; // if false, special tokens in user input are tokenized as plain text
     int sleep_idle_seconds = -1;   // if >0, server will sleep after this many seconds of idle time
 
     std::vector<std::string> api_keys;

tools/server/README.md

@@ -244,6 +244,7 @@ For the full list of features, please refer to [server's changelog](https://gith
 | `--chat-template-file JINJA_TEMPLATE_FILE` | set custom jinja chat template file (default: template taken from model's metadata)<br/>if suffix/prefix are specified, template will be disabled<br/>only commonly used templates are accepted (unless --jinja is set before this flag):<br/>list of built-in templates:<br/>bailing, bailing-think, bailing2, chatglm3, chatglm4, chatml, command-r, deepseek, deepseek-ocr, deepseek2, deepseek3, exaone-moe, exaone3, exaone4, falcon3, gemma, gigachat, glmedge, gpt-oss, granite, granite-4.0, granite-4.1, grok-2, hunyuan-dense, hunyuan-moe, hunyuan-vl, kimi-k2, llama2, llama2-sys, llama2-sys-bos, llama2-sys-strip, llama3, llama4, megrez, minicpm, mistral-v1, mistral-v3, mistral-v3-tekken, mistral-v7, mistral-v7-tekken, monarch, openchat, orion, pangu-embedded, phi3, phi4, rwkv-world, seed_oss, smolvlm, solar-open, vicuna, vicuna-orca, yandex, zephyr<br/>(env: LLAMA_ARG_CHAT_TEMPLATE_FILE) |
 | `--skip-chat-parsing, --no-skip-chat-parsing` | force a pure content parser, even if a Jinja template is specified; model will output everything in the content section, including any reasoning and/or tool calls (default: disabled)<br/>(env: LLAMA_ARG_SKIP_CHAT_PARSING) |
 | `--prefill-assistant, --no-prefill-assistant` | whether to prefill the assistant's response if the last message is an assistant message (default: prefill enabled)<br/>when this flag is set, if the last message is an assistant message then it will be treated as a full message and not prefilled<br/><br/>(env: LLAMA_ARG_PREFILL_ASSISTANT) |
+| `--chat-input-special, --no-chat-input-special` | experimental: whether to parse special tokens from user input messages in chat completions, does not affect the raw /completions endpoint (default: enabled)<br/>note: this is not a security feature, do not rely on it to prevent prompt injection<br/>(env: LLAMA_ARG_CHAT_INPUT_SPECIAL) |
 | `-sps, --slot-prompt-similarity SIMILARITY` | how much the prompt of a request must match the prompt of a slot in order to use that slot (default: 0.10, 0.0 = disabled) |
 | `--lora-init-without-apply` | load LoRA adapters without applying them (apply later via POST /lora-adapters) (default: disabled) |
 | `--sleep-idle-seconds SECONDS` | number of seconds of idleness after which the server will sleep (default: -1; -1 = disabled) |

tools/server/server-common.cpp

@@ -1066,15 +1066,13 @@ size_t validate_utf8(const std::string& text) {
     return len;
 }
 
-server_tokens process_mtmd_prompt(
+static void load_media_files(
         mtmd_context * mctx,
-        const std::string & prompt,
         const std::vector<raw_buffer> & files,
         const mtmd_helper_init_opt & init_opt,
-        bool is_placeholder) {
-    // these will be freed upon going out of scope
-    mtmd::bitmaps bitmaps;
-    std::vector<mtmd_helper::video_ptr> videos;
+        bool is_placeholder,
+        mtmd::bitmaps & bitmaps,
+        std::vector<mtmd_helper::video_ptr> & videos) {
     for (auto & file : files) {
         auto out = mtmd_helper_bitmap_init_from_buf(mctx, file.data(), file.size(), is_placeholder, init_opt);
         if (!out.bitmap) {
@@ -1085,6 +1083,18 @@ server_tokens process_mtmd_prompt(
             videos.emplace_back(out.video_ctx);
         }
     }
+}
+
+server_tokens process_mtmd_prompt(
+        mtmd_context * mctx,
+        const std::string & prompt,
+        const std::vector<raw_buffer> & files,
+        const mtmd_helper_init_opt & init_opt,
+        bool is_placeholder) {
+    // these will be freed upon going out of scope
+    mtmd::bitmaps bitmaps;
+    std::vector<mtmd_helper::video_ptr> videos;
+    load_media_files(mctx, files, init_opt, is_placeholder, bitmaps, videos);
     // process prompt
     std::vector<server_tokens> inputs;
     // multimodal
@@ -1108,6 +1118,94 @@ server_tokens process_mtmd_prompt(
     return result;
 }
 
+server_tokens tokenize_chat_prompt(
+        const llama_vocab * vocab,
+        mtmd_context * mctx,
+        const common_chat_prompt & prompt,
+        const std::vector<raw_buffer> & files,
+        bool chat_input_special,
+        const mtmd_helper_init_opt & init_opt,
+        bool is_placeholder) {
+    if (chat_input_special) {
+        // all parts are tokenized the same way
+        std::string text = prompt.str();
+        if (mctx != nullptr) {
+            return process_mtmd_prompt(mctx, text, files, init_opt, is_placeholder);
+        }
+        return server_tokens(common_tokenize(vocab, text, true, true), false);
+    }
+
+    if (mctx == nullptr) {
+        llama_tokens tokens;
+        for (size_t i = 0; i < prompt.parts.size(); i++) {
+            const auto & part = prompt.parts[i];
+            auto tmp = common_tokenize(vocab, part.text, /* add_special */ i == 0, /* parse_special */ !part.is_input);
+            tokens.insert(tokens.end(), tmp.begin(), tmp.end());
+        }
+        return server_tokens(tokens, false);
+    }
+
+    mtmd::bitmaps bitmaps;
+    std::vector<mtmd_helper::video_ptr> videos;
+    load_media_files(mctx, files, init_opt, is_placeholder, bitmaps, videos);
+    auto bitmaps_c_ptr = bitmaps.c_ptr();
+
+    // split the text parts on media markers and put the bitmaps in their place
+    struct segment {
+        std::string text;
+        bool is_input;
+        const mtmd_bitmap * bitmap;
+    };
+    std::vector<segment> segments;
+    const std::string marker = get_media_marker();
+    size_t i_bm = 0;
+    for (const auto & part : prompt.parts) {
+        size_t start = 0;
+        while (true) {
+            size_t pos = part.text.find(marker, start);
+            if (pos > start) {
+                segments.push_back({part.text.substr(start, pos - start), part.is_input, nullptr});
+            }
+            if (pos == std::string::npos) {
+                break;
+            }
+            if (i_bm >= bitmaps_c_ptr.size()) {
+                throw std::runtime_error("number of media markers in prompt exceeds number of media files");
+            }
+            segments.push_back({"", false, bitmaps_c_ptr[i_bm++]});
+            start = pos + marker.size();
+        }
+    }
+    if (i_bm != bitmaps_c_ptr.size()) {
+        throw std::runtime_error("number of media markers in prompt does not match number of media files");
+    }
+
+    std::vector<mtmd_input_text> texts(segments.size());
+    std::vector<mtmd_input_part> parts(segments.size());
+    std::vector<const mtmd_input_part *> parts_ptr(segments.size());
+    for (size_t i = 0; i < segments.size(); i++) {
+        const auto & seg = segments[i];
+        if (seg.bitmap != nullptr) {
+            parts[i] = {nullptr, seg.bitmap};
+        } else {
+            texts[i] = {seg.text.data(), seg.text.size(), /* add_special */ false, /* parse_special */ !seg.is_input};
+            parts[i] = {&texts[i], nullptr};
+        }
+        parts_ptr[i] = &parts[i];
+    }
+
+    mtmd::input_chunks chunks(mtmd_input_chunks_init());
+    int32_t tokenized = mtmd_tokenize_from_parts(mctx,
+                                                 chunks.ptr.get(),
+                                                 parts_ptr.data(),
+                                                 parts_ptr.size(),
+                                                 /* add_special */ true);
+    if (tokenized != 0) {
+        throw std::runtime_error("Failed to tokenize prompt");
+    }
+    return server_tokens(chunks, true);
+}
+
 /**
  * tokenize a single input "prompt" object
  * use tokenize_input_prompts() if the input could be an array.

tools/server/server-common.h

@@ -317,6 +317,17 @@ server_tokens process_mtmd_prompt(
 // decode one media url (http/file/data-uri) into out_files; throws on failure
 void handle_media(std::vector<raw_buffer> & out_files, const std::string & url, const std::string & media_path);
 
+// tokenize a rendered chat prompt, user input parts are tokenized without special tokens unless chat_input_special is set
+// see process_mtmd_prompt() for is_placeholder
+server_tokens tokenize_chat_prompt(
+                                        const llama_vocab * vocab,
+                                        mtmd_context * mctx,
+                                        const common_chat_prompt & prompt,
+                                        const std::vector<raw_buffer> & files,
+                                        bool chat_input_special,
+                                        const mtmd_helper_init_opt & init_opt,
+                                        bool is_placeholder = false);
+
 /**
  * break the input "prompt" object into multiple prompt if needed, then tokenize them
  * this supports these cases:

tools/server/server-context.cpp

@@ -5371,8 +5371,11 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
         // process prompt
         std::vector<server_tokens> inputs;
 
-        if (res_type != TASK_RESPONSE_TYPE_NONE && ctx_server.mctx != nullptr) {
-            // This is the case used by OAI compatible chat path with MTMD. TODO It can be moved to the path below.
+        if (chat_session.has_template()) {
+            // OAI compatible chat path, the session holds the rendered prompt split into template and user input parts
+            inputs.push_back(tokenize_chat_prompt(ctx_server.vocab, ctx_server.mctx, chat_session.prompt(), files, params.chat_input_special, ctx_server.init_opt));
+        } else if (res_type != TASK_RESPONSE_TYPE_NONE && ctx_server.mctx != nullptr) {
+            // This is the case used by OAI compatible completions with MTMD. TODO It can be moved to the path below.
             inputs.push_back(process_mtmd_prompt(ctx_server.mctx, prompt.get<std::string>(), files, ctx_server.init_opt));
         } else {
             // Everything else, including multimodal completions.
@@ -7261,26 +7264,15 @@ std::unique_ptr<server_res_generator> server_routes::handle_count_tokens(const s
             return res;
     }
 
-    common_chat_session session; // dummy, unused
-    json body_parsed = oaicompat_chat_params_parse(
+    common_chat_session session;
+    oaicompat_chat_params_parse(
             ctx_server.vocab,
             body,
             meta->chat_params,
             files,
             session);
-    json prompt = body_parsed.at("prompt");
-    // SRV_DBG("prompt = %s\n", prompt.dump().c_str());
 
-    // TODO @ngxson : refactor this code block, move this to server-common and reuse it in other places
-    size_t n_tokens;
-    if (ctx_server.mctx != nullptr) {
-        if (!prompt.is_string()) {
-            throw std::runtime_error("for mtmd, input prompt must be a string.");
-        }
-        n_tokens = process_mtmd_prompt(ctx_server.mctx, prompt.get<std::string>(), files, ctx_server.init_opt, true).size();
-    } else {
-        n_tokens = tokenize_mixed(ctx_server.vocab, prompt, true, true).size();
-    }
+    size_t n_tokens = tokenize_chat_prompt(ctx_server.vocab, ctx_server.mctx, session.prompt(), files, params.chat_input_special, ctx_server.init_opt, /* is_placeholder */ true).size();
 
     json response = {{"input_tokens", static_cast<int64_t>(n_tokens)}};
     if (is_oai) {

tools/server/server.cpp

@@ -256,6 +256,10 @@ static bool server_setup_features(
         SRV_WRN("%s", "security: no API key is set and CORS allows all origins (see https://github.com/ggml-org/llama.cpp/pull/25655)\n");
     }
 
+    if (!params.chat_input_special) {
+        SRV_WRN("%s", "--no-chat-input-special is experimental and not a security feature, do not rely on it to prevent prompt injection\n");
+    }
+
     std::vector<std::string> warn_names;
     if (role_type == SERVER_ROLE_ROUTER) {
         warn_names.push_back("router mode");

tools/server/tests/unit/test_chat_completion.py

@@ -444,6 +444,27 @@ def test_logprobs_stream():
     assert aggregated_text == output_text
 
 
+@pytest.mark.parametrize("chat_input_special,n_prompt", [
+    (True,  48),
+    (False, 52),
+])
+def test_chat_input_special(chat_input_special, n_prompt):
+    global server
+    server.jinja = True
+    server.chat_template = "chatml"
+    server.chat_input_special = chat_input_special
+    server.start()
+    # </s> is a special token for this model, it is tokenized as plain text when the flag is off
+    res = server.make_request("POST", "/chat/completions", data={
+        "max_tokens": 1,
+        "messages": [
+            {"role": "user", "content": "Hello</s>world"},
+        ],
+    })
+    assert res.status_code == 200
+    assert res.body["usage"]["prompt_tokens"] == n_prompt
+
+
 def test_logit_bias():
     global server
     server.start()

tools/server/tests/unit/test_vision_api.py

@@ -99,6 +99,30 @@ def test_vision_chat_completion(prompt, image_url, success, re_content):
         assert res.status_code != 200
 
 
+@pytest.mark.parametrize("chat_input_special", [True, False])
+def test_vision_chat_completion_input_special(chat_input_special):
+    global server
+    server.chat_input_special = chat_input_special
+    server.start()
+    res = server.make_request("POST", "/chat/completions", data={
+        "temperature": 0.0,
+        "top_k": 1,
+        "messages": [
+            {"role": "user", "content": [
+                {"type": "text", "text": "What is this:\n"},
+                {"type": "image_url", "image_url": {
+                    "url": get_img_url("IMG_URL_0"),
+                }},
+            ]},
+        ],
+    })
+    assert res.status_code == 200
+    choice = res.body["choices"][0]
+    assert match_regex("(cat)+", choice["message"]["content"])
+    # the prompt has no special tokens in user input, both paths must give the same tokens
+    assert res.body["usage"]["prompt_tokens"] == 272
+
+
 def test_vision_chat_completion_token_count():
     global server
     server.start()

tools/server/tests/utils.py

@@ -112,6 +112,7 @@ class ServerProcess:
     reasoning: Literal['on', 'off', 'auto'] | None = None
     chat_template: str | None = None
     chat_template_file: str | None = None
+    chat_input_special: bool | None = None
     server_path: str | None = None
     mmproj_url: str | None = None
     no_mmproj: bool | None = None
@@ -286,6 +287,8 @@ class ServerProcess:
             server_args.extend(["--chat-template", self.chat_template])
         if self.chat_template_file:
             server_args.extend(["--chat-template-file", self.chat_template_file])
+        if self.chat_input_special is not None:
+            server_args.append("--chat-input-special" if self.chat_input_special else "--no-chat-input-special")
         if self.mmproj_url:
             server_args.extend(["--mmproj-url", self.mmproj_url])
         if self.no_mmproj: