From e587faf5fe2c96f666a2f89b156904516fb11d95 Mon Sep 17 00:00:00 2001 From: Callum Mitchell Date: Tue, 8 Sep 2026 15:22:50 -0600 Subject: [PATCH 1/2] Fix: comprehensive token ID validation with 'invalid_prompt' errors MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit POST /v1/completions requests with integer-array prompts now correctly reject invalid token IDs with HTTP 400 and error.type "invalid_prompt": - Negative token IDs (except -1) - -1 (LLAMA_TOKEN_NULL) without media chunks - Out-of-vocabulary token IDs - Non-integer types: floats, booleans, nulls - Integer values outside int32 range (int64 overflow) The fix validates token types BEFORE parsing to prevent implicit JSON conversions (float truncation, bool->int, int64 wrapping). Code changes: - json_get_token() helper for type-safe token extraction - Throw std::invalid_argument for validation errors - Three-tier exception handling in /v1/completions endpoint: 1. Prompt extraction → invalid_request_error 2. Tokenization → invalid_prompt 3. Schema validation & task creation → invalid_request_error - Task type check in launch_slot_with_task: completions use invalid_prompt, other endpoints use invalid_request_error - Other endpoints rely on global exception handler for type errors Test coverage: new unit tests across test_completion.py, test_embedding.py, and test_tokenize.py covering token validation and error type checks. Co-Authored-By: Claude Signed-off-by: Callum Mitchell --- tools/server/server-common.cpp | 32 ++++++++++-- tools/server/server-common.h | 1 + tools/server/server-context.cpp | 27 +++++++--- tools/server/tests/unit/test_completion.py | 59 ++++++++++++++++++++++ tools/server/tests/unit/test_embedding.py | 59 +++++++++++++++++++++- tools/server/tests/unit/test_tokenize.py | 16 ++++++ 6 files changed, 181 insertions(+), 13 deletions(-) diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp index c32c0d1a5aa7..323c4a95fc05 100644 --- a/tools/server/server-common.cpp +++ b/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(); + if (val < INT32_MIN || val > INT32_MAX) { + throw std::invalid_argument("Token ID out of int32 range: " + std::to_string(val)); + } + return static_cast(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()); + 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 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 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 { diff --git a/tools/server/server-common.h b/tools/server/server-common.h index 5bd2ed57ab89..6703abfff20f 100644 --- a/tools/server/server-common.h +++ b/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 diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 60c563345399..f34628ea123f 100644 --- a/tools/server/server-context.cpp +++ b/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_routes::handle_completions_impl( int32_t sse_ping_interval = params.sse_ping_interval; - try { - std::vector tasks; + std::vector 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().c_str() : prompt.dump(2).c_str()); @@ -4814,10 +4818,16 @@ std::unique_ptr 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 inputs; + // process prompt + std::vector 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(), files, ctx_server.init_opt)); @@ -4825,9 +4835,14 @@ std::unique_ptr server_routes::handle_completions_impl( // 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); diff --git a/tools/server/tests/unit/test_completion.py b/tools/server/tests/unit/test_completion.py index b62abc6614ff..facac6007a7b 100644 --- a/tools/server/tests/unit/test_completion.py +++ b/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 diff --git a/tools/server/tests/unit/test_embedding.py b/tools/server/tests/unit/test_embedding.py index 9efa273fc9fe..c2df74b4ad33 100644 --- a/tools/server/tests/unit/test_embedding.py +++ b/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" diff --git a/tools/server/tests/unit/test_tokenize.py b/tools/server/tests/unit/test_tokenize.py index 424cac5f3d39..5e276e061371 100644 --- a/tools/server/tests/unit/test_tokenize.py +++ b/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" From ed405de7cfffc06a1266d51e954bef748368ac89 Mon Sep 17 00:00:00 2001 From: Callum Mitchell Date: Tue, 27 Oct 2026 12:11:08 -0600 Subject: [PATCH 2/2] Token ID validation: properly handle uint64 max case + media loading Addresses missing coverage: - Multimodal exception classification: added runtime_error catch block in completions endpoint to return 400 invalid_request_error for media loading/tokenization failures instead of 500. - Unsigned int64 max: added unsigned bounds check in json_get_token() before signed conversion to prevent uint64 max (18446744073709551615) from wrapping to -1. - Added is_number_unsigned() to common_json wrapper - Token ID valudation test in each endpoint (completions, embeddings, tokenize) to verify proper uint64 max value handling Co-authored-by: Claude Signed-off-by: Callum Mitchell --- common/json.cpp | 17 +++++++++-------- common/json.h | 1 + tools/server/server-common.cpp | 9 +++++++++ tools/server/server-context.cpp | 4 ++++ tools/server/tests/unit/test_completion.py | 6 ++++++ tools/server/tests/unit/test_embedding.py | 5 +++++ tools/server/tests/unit/test_tokenize.py | 5 +++++ 7 files changed, 39 insertions(+), 8 deletions(-) diff --git a/common/json.cpp b/common/json.cpp index 37713cef29e7..cf49845c3f3b 100644 --- a/common/json.cpp +++ b/common/json.cpp @@ -251,14 +251,15 @@ common_json common_json::make(const common_json_value & val) { return common_json(val); } -bool common_json::is_null() const { return as_json(this).is_null(); } -bool common_json::is_object() const { return as_json(this).is_object(); } -bool common_json::is_array() const { return as_json(this).is_array(); } -bool common_json::is_string() const { return as_json(this).is_string(); } -bool common_json::is_boolean() const { return as_json(this).is_boolean(); } -bool common_json::is_number() const { return as_json(this).is_number(); } -bool common_json::is_number_integer() const { return as_json(this).is_number_integer(); } -bool common_json::is_number_float() const { return as_json(this).is_number_float(); } +bool common_json::is_null() const { return as_json(this).is_null(); } +bool common_json::is_object() const { return as_json(this).is_object(); } +bool common_json::is_array() const { return as_json(this).is_array(); } +bool common_json::is_string() const { return as_json(this).is_string(); } +bool common_json::is_boolean() const { return as_json(this).is_boolean(); } +bool common_json::is_number() const { return as_json(this).is_number(); } +bool common_json::is_number_integer() const { return as_json(this).is_number_integer(); } +bool common_json::is_number_unsigned() const { return as_json(this).is_number_unsigned(); } +bool common_json::is_number_float() const { return as_json(this).is_number_float(); } bool common_json::empty() const { return as_json(this).empty(); } size_t common_json::size() const { return as_json(this).size(); } diff --git a/common/json.h b/common/json.h index 95b6d6f212e5..c6b69aa602ee 100644 --- a/common/json.h +++ b/common/json.h @@ -183,6 +183,7 @@ class common_json { bool is_boolean() const; bool is_number() const; bool is_number_integer() const; + bool is_number_unsigned() const; bool is_number_float() const; bool empty() const; diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp index 323c4a95fc05..a9122271cd9e 100644 --- a/tools/server/server-common.cpp +++ b/tools/server/server-common.cpp @@ -870,6 +870,15 @@ static llama_token json_get_token(const json & elem) { if (!elem.is_number_integer()) { throw std::invalid_argument("Token ID must be an integer"); } + + // Check uint64 range first to prevent false validation due to int overflow + if (elem.is_number_unsigned()) { + uint64_t uval = elem.get(); + if (uval > static_cast(INT32_MAX)) { + throw std::invalid_argument("Token ID out of int32 range: " + std::to_string(uval)); + } + } + int64_t val = elem.get(); if (val < INT32_MIN || val > INT32_MAX) { throw std::invalid_argument("Token ID out of int32 range: " + std::to_string(val)); diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index f34628ea123f..00d50fab7ca2 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -4838,6 +4838,10 @@ std::unique_ptr server_routes::handle_completions_impl( } catch (const std::invalid_argument & e) { res->error(format_error_response(e.what(), ERROR_TYPE_INVALID_PROMPT)); return res; + } catch (const std::runtime_error & e) { + // Multimodal loading/tokenization failures + res->error(format_error_response(e.what(), ERROR_TYPE_INVALID_REQUEST)); + return res; } // tasks.reserve(inputs.size()); // TODO: this is inaccurate due to child tasks diff --git a/tools/server/tests/unit/test_completion.py b/tools/server/tests/unit/test_completion.py index facac6007a7b..4dd4c5e90b9d 100644 --- a/tools/server/tests/unit/test_completion.py +++ b/tools/server/tests/unit/test_completion.py @@ -714,6 +714,8 @@ def test_completion_prompt_cache(): ([1, "x", None], 400, "invalid_prompt"), # Invalid: int32 overflow ([1, 4294967297], 400, "invalid_prompt"), + # Invalid: uint64 max - must not overflow to -1 (valid) + ([18446744073709551615], 400, "invalid_prompt"), ]) def test_token_validation(prompt, expected_status, expected_error_type): """ @@ -731,6 +733,10 @@ def test_token_validation(prompt, expected_status, expected_error_type): if expected_error_type: assert "error" in res.body assert res.body["error"]["type"] == expected_error_type + # uint64 max: verify error message to confirm proper bounds checking + # vs. silent wrap-around to -1 (LLAMA_TOKEN_NULL) which would trigger "Invalid token IDs" + if 18446744073709551615 in prompt: + assert "Token ID out of int32 range: 18446744073709551615" in res.body["error"]["message"] @pytest.mark.parametrize("data,expected_error_type", [ diff --git a/tools/server/tests/unit/test_embedding.py b/tools/server/tests/unit/test_embedding.py index c2df74b4ad33..cbbc8dea1c99 100644 --- a/tools/server/tests/unit/test_embedding.py +++ b/tools/server/tests/unit/test_embedding.py @@ -366,6 +366,7 @@ def test_embedding_invalid_token_ids(input): [1, -5], [1, 99999999], [1, -1], # Token -1 without media + [18446744073709551615], # uint64 max - must not overflow to -1 (valid) ]) def test_embedding_invalid_token_range(input): """Negative and out-of-vocab tokens should return 400 invalid_request_error.""" @@ -375,6 +376,10 @@ def test_embedding_invalid_token_range(input): res = server.make_request("POST", "/v1/embeddings", data={"input": input}) assert res.status_code == 400 assert res.body["error"]["type"] == "invalid_request_error" + # uint64 max: verify error message to confirm proper bounds checking + # vs. silent wrap-around to -1 (LLAMA_TOKEN_NULL) which would trigger "Invalid token IDs" + if 18446744073709551615 in input: + assert "Token ID out of int32 range: 18446744073709551615" in res.body["error"]["message"] @pytest.mark.parametrize("input", [ diff --git a/tools/server/tests/unit/test_tokenize.py b/tools/server/tests/unit/test_tokenize.py index 5e276e061371..a4acae1671ba 100644 --- a/tools/server/tests/unit/test_tokenize.py +++ b/tools/server/tests/unit/test_tokenize.py @@ -65,6 +65,7 @@ def test_tokenize_with_pieces(): [False], [1, "x", None], [1, 4294967297], + [18446744073709551615], # uint64 max - must not overflow to -1 (valid) ]) def test_tokenize_invalid_token_ids(content): """Invalid token IDs in mixed arrays should return 400.""" @@ -73,3 +74,7 @@ def test_tokenize_invalid_token_ids(content): res = server.make_request("POST", "/tokenize", data={"content": content}) assert res.status_code == 400 assert res.body["error"]["type"] == "invalid_request_error" + # uint64 max: verify error message to confirm proper bounds checking + # vs. silent wrap-around to -1 which would return 200 success with [-1] tokens + if 18446744073709551615 in content: + assert "Token ID out of int32 range: 18446744073709551615" in res.body["error"]["message"]