Fix: Return 'invalid_prompt' error type for invalid token IDs

tools/server/server-common.cpp

@@ -69,6 +69,10 @@ json format_error_response(const std::string & message, const enum error_type ty
             type_str = "exceed_context_size_error";
             code = 400;
             break;
+        case ERROR_TYPE_INVALID_PROMPT:
+            type_str = "invalid_prompt";
+            code = 400;
+            break;
     }
     return json {
         {"code", code},
@@ -862,6 +866,17 @@ bool json_is_array_of_numbers(const json & data) {
     return false;
 }
 
+static llama_token json_get_token(const json & elem) {
+    if (!elem.is_number_integer()) {
+        throw std::invalid_argument("Token ID must be an integer");
+    }
+    int64_t val = elem.get<int64_t>();
+    if (val < INT32_MIN || val > INT32_MAX) {
+        throw std::invalid_argument("Token ID out of int32 range: " + std::to_string(val));
+    }
+    return static_cast<llama_token>(val);
+}
+
 bool json_is_array_of_mixed_numbers_strings(const json & data) {
     bool seen_string = false;
     bool seen_number = false;
@@ -935,7 +950,7 @@ llama_tokens tokenize_mixed(const llama_vocab * vocab, const json & json_prompt,
                     first = false;
                 }
 
-                prompt_tokens.push_back(p.template get<llama_token>());
+                prompt_tokens.push_back(json_get_token(p));
             }
         }
     } else {
@@ -1033,14 +1048,17 @@ server_tokens tokenize_input_subprompt(const llama_vocab * vocab, mtmd_context *
         llama_tokens tmp = tokenize_mixed(vocab, json_prompt, add_special, parse_special);
         return server_tokens(tmp, false);
     } else if (json_is_array_of_numbers(json_prompt)) {
-        // array of tokens
-        llama_tokens tmp = json_prompt.get<llama_tokens>();
+        llama_tokens tmp;
+        tmp.reserve(json_prompt.size());
+        for (const auto & elem : json_prompt) {
+            tmp.push_back(json_get_token(elem));
+        }
         return server_tokens(tmp, false);
     } else if (json_prompt.contains(JSON_STRING_PROMPT_KEY)) {
         // JSON object with prompt key.
         if (json_prompt.contains(JSON_MTMD_DATA_KEY)) {
             if (!has_mtmd)
-                throw std::runtime_error("Multimodal data provided, but model does not support multimodal requests.");
+                throw std::invalid_argument("Multimodal data provided, but model does not support multimodal requests.");
 
             // JSON object with prompt and multimodal key.
             std::vector<raw_buffer> files;
@@ -1050,7 +1068,11 @@ server_tokens tokenize_input_subprompt(const llama_vocab * vocab, mtmd_context *
             return process_mtmd_prompt(mctx, json_prompt.at(JSON_STRING_PROMPT_KEY), files, init_opt);
         } else {
             // Not multimodal, but contains a subobject.
-            llama_tokens tmp = tokenize_mixed(vocab, json_prompt.at(JSON_STRING_PROMPT_KEY), add_special, parse_special);
+            const auto & prompt_str = json_prompt.at(JSON_STRING_PROMPT_KEY);
+            if (!prompt_str.is_string() && !prompt_str.is_array()) {
+                throw std::invalid_argument("\"prompt_string\" must be a string or array of tokens");
+            }
+            llama_tokens tmp = tokenize_mixed(vocab, prompt_str, add_special, parse_special);
             return server_tokens(tmp, false);
         }
    } else {

tools/server/server-common.h

@@ -64,6 +64,7 @@ enum error_type {
     ERROR_TYPE_UNAVAILABLE, // custom error
     ERROR_TYPE_NOT_SUPPORTED, // custom error
     ERROR_TYPE_EXCEED_CONTEXT_SIZE, // custom error
+    ERROR_TYPE_INVALID_PROMPT, // custom error - for bad token IDs in prompt array
 };
 
 // thin wrapper around common_grammar_trigger with (de)serialization functions

tools/server/server-context.cpp

@@ -1917,7 +1917,10 @@ struct server_context_impl {
         }
 
         if (!task.tokens.validate(ctx_tgt)) {
-            send_error(task, "Prompt contains invalid tokens", ERROR_TYPE_INVALID_REQUEST);
+            error_type err_type = (task.type == SERVER_TASK_TYPE_COMPLETION)
+                ? ERROR_TYPE_INVALID_PROMPT
+                : ERROR_TYPE_INVALID_REQUEST;
+            send_error(task, "Invalid token IDs", err_type);
             return false;
         }
 
@@ -4798,9 +4801,10 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
 
     int32_t sse_ping_interval = params.sse_ping_interval;
 
-    try {
-        std::vector<server_task> tasks;
+    std::vector<server_task> tasks;
 
+    // Extract prompt
+    try {
         const auto & prompt = data.at("prompt");
         // TODO: this log can become very long, put it behind a flag or think about a more compact format
         //SRV_DBG("Prompt: %s\n", prompt.is_string() ? prompt.get<std::string>().c_str() : prompt.dump(2).c_str());
@@ -4814,20 +4818,31 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
                 SRV_ERR("failed to create %s\n", file_path.string().c_str());
             }
         }
+    } catch (const std::exception & e) {
+        res->error(format_error_response(e.what(), ERROR_TYPE_INVALID_REQUEST));
+        return res;
+    }
 
-        // process prompt
-        std::vector<server_tokens> inputs;
+    // process prompt
+    std::vector<server_tokens> inputs;
 
+    try {
+        const auto & prompt = data.at("prompt");
         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.
             inputs.push_back(process_mtmd_prompt(ctx_server.mctx, prompt.get<std::string>(), files, ctx_server.init_opt));
         } else {
             // Everything else, including multimodal completions.
             inputs = tokenize_input_prompts(ctx_server.vocab, ctx_server.mctx, prompt, true, true, ctx_server.init_opt);
         }
+    } catch (const std::invalid_argument & e) {
+        res->error(format_error_response(e.what(), ERROR_TYPE_INVALID_PROMPT));
+        return res;
+    }
 
-        // tasks.reserve(inputs.size()); // TODO: this is inaccurate due to child tasks
+    // tasks.reserve(inputs.size()); // TODO: this is inaccurate due to child tasks
 
+    try {
         for (size_t i = 0; i < inputs.size(); i++) {
             server_task task = server_task(type);
 

tools/server/tests/unit/test_completion.py

@@ -689,3 +689,62 @@ def test_completion_prompt_cache():
         assert "prompt_n" in timings and timings["prompt_n"] + timings["cache_n"] == n_prompt
         assert "predicted_n" in timings and timings["predicted_n"] == n_predict
         assert "tokens" in res.body and isinstance(res.body["tokens"], list)
+
+
+@pytest.mark.parametrize("prompt,expected_status,expected_error_type", [
+    # Valid
+    ([1], 200, None),
+    ([1, 2, 3], 200, None),
+    ([1, "hello", 2], 200, None),
+    (["HeLlO 世界! 🎉, [Thi$ shou1d_be-\"valid\'..}"], 200, None),
+    # Invalid: negative (except -1)
+    ([1, -5], 400, "invalid_prompt"),
+    # Invalid: token ID out of vocabulary
+    ([99999999], 400, "invalid_prompt"),
+    # Invalid: LLAMA_TOKEN_NULL (-1) without media chunk
+    ([1, -1], 400, "invalid_prompt"),
+    # Invalid: float
+    ([2.5, 1], 400, "invalid_prompt"),
+    ([1, "x", 2.0], 400, "invalid_prompt"),
+    # Invalid: bool
+    ([False], 400, "invalid_prompt"),
+    ([1, "x", True], 400, "invalid_prompt"),
+    # Invalid: null
+    ([None], 400, "invalid_prompt"),
+    ([1, "x", None], 400, "invalid_prompt"),
+    # Invalid: int32 overflow
+    ([1, 4294967297], 400, "invalid_prompt"),
+])
+def test_token_validation(prompt, expected_status, expected_error_type):
+    """
+    POST /v1/completions with integer-array prompt must reject negative,
+    out-of-vocabulary, or non-integer token IDs with HTTP 400 and
+    error.type "invalid_prompt".
+    """
+    global server
+    server.start()
+    res = server.make_request("POST", "/v1/completions", data={
+        "prompt": prompt,
+        "max_tokens": 1,
+    })
+    assert res.status_code == expected_status
+    if expected_error_type:
+        assert "error" in res.body
+        assert res.body["error"]["type"] == expected_error_type
+
+
+@pytest.mark.parametrize("data,expected_error_type", [
+    ({"prompt": "hello", "temperature": "hot"}, "invalid_request_error"),
+    ({"prompt": "hello", "top_k": "many"}, "invalid_request_error"),
+    ({}, "invalid_request_error"),
+    ({"prompt": []}, "invalid_prompt"),  # Empty array
+    ({"prompt": {"prompt_string": 123}}, "invalid_prompt"),  # Non-string prompt_string
+    ({"prompt": {"prompt_string": "test", "multimodal_data": []}}, "invalid_prompt"),  # Multimodal on non-multimodal model
+])
+def test_invalid_request_params(data, expected_error_type):
+    """Invalid request parameters should return 400 with correct error type."""
+    global server
+    server.start()
+    res = server.make_request("POST", "/v1/completions", data=data)
+    assert res.status_code == 400
+    assert res.body["error"]["type"] == expected_error_type

tools/server/tests/unit/test_embedding.py

@@ -123,7 +123,6 @@ def test_embedding_content_text_same_as_string():
 @pytest.mark.parametrize(
     "input",
     [
-        [],
         {"content": "string"},
         {"content": [{"type": "unknown"}]},
         # model is not multimodal
@@ -332,7 +331,6 @@ def test_embedding_openai_library_base64():
 @pytest.mark.parametrize(
     "data",
     [
-        {"input": []},
         {"input": True},
         {"input": "hello", "encoding_format": 1},
     ]
@@ -345,3 +343,60 @@ def test_embedding_invalid_request(data):
     assert res.status_code == 400
     assert "error" in res.body
     assert res.body["error"]["type"] == "invalid_request_error"
+
+
+@pytest.mark.parametrize("input", [
+    [1, "x", 2.5],
+    [1, True],
+    [False],
+    [1, "x", None],
+    [1, 4294967297],
+])
+def test_embedding_invalid_token_ids(input):
+    """Invalid token IDs in input arrays should return 400."""
+    global server
+    server.pooling = 'last'
+    server.start()
+    res = server.make_request("POST", "/v1/embeddings", data={"input": input})
+    assert res.status_code == 400
+    assert res.body["error"]["type"] == "invalid_request_error"
+
+
+@pytest.mark.parametrize("input", [
+    [1, -5],
+    [1, 99999999],
+    [1, -1],  # Token -1 without media
+])
+def test_embedding_invalid_token_range(input):
+    """Negative and out-of-vocab tokens should return 400 invalid_request_error."""
+    global server
+    server.pooling = 'last'
+    server.start()
+    res = server.make_request("POST", "/v1/embeddings", data={"input": input})
+    assert res.status_code == 400
+    assert res.body["error"]["type"] == "invalid_request_error"
+
+
+@pytest.mark.parametrize("input", [
+    [],  # Empty array
+    {"prompt_string": 123},  # Non-string prompt_string
+    {"prompt_string": "test", "multimodal_data": []},  # Multimodal on non-multimodal model
+])
+def test_embedding_invalid_prompt_format(input):
+    """Invalid prompt formats should return 400 invalid_request_error."""
+    global server
+    server.pooling = 'last'
+    server.start()
+    res = server.make_request("POST", "/v1/embeddings", data={"input": input})
+    assert res.status_code == 400
+    assert res.body["error"]["type"] == "invalid_request_error"
+
+
+def test_embedding_missing_input():
+    """Missing input field should return 400 invalid_request_error."""
+    global server
+    server.pooling = 'last'
+    server.start()
+    res = server.make_request("POST", "/v1/embeddings", data={})
+    assert res.status_code == 400
+    assert res.body["error"]["type"] == "invalid_request_error"

tools/server/tests/unit/test_tokenize.py

@@ -57,3 +57,19 @@ def test_tokenize_with_pieces():
         assert token["id"] > 0
         assert "piece" in token
         assert len(token["piece"]) > 0
+
+
+@pytest.mark.parametrize("content", [
+    [1, "x", 2.5],
+    [1, True],
+    [False],
+    [1, "x", None],
+    [1, 4294967297],
+])
+def test_tokenize_invalid_token_ids(content):
+    """Invalid token IDs in mixed arrays should return 400."""
+    global server
+    server.start()
+    res = server.make_request("POST", "/tokenize", data={"content": content})
+    assert res.status_code == 400
+    assert res.body["error"]["type"] == "invalid_request_error"