koboldcpp/tests/test-server-component.cpp
Yuri Khrustalev 04f7226a2c
server : escape the special tokens without std::regex (#30243)
The d1-omni and kev prompts escape "<|name|>" in the request text with
std::regex_replace. libstdc++ recurses for each char of the name, so a
state of "<|" and 30000 letters overflows the stack of the HTTP thread
and llama-server dies with SIGSEGV.

server_decision_escape_special_tokens() writes the same text with a loop.
test-server-component compares it with the regex on every string of up to
6 chars of "<|>a-".
2026-10-14 01:17:11 +02:00

459 lines
18 KiB
C++

#include "testing.h"
#include "../tools/server/server-common.h"
#include "../tools/server/server-decision.h"
#include <cstdlib>
#include <regex>
#include <string>
#include <thread>
#include <utility>
#include <vector>
// unit tests for server internal components, no model required
// NOTE: end-to-end tests live in tools/server/tests
// see skills/code-review/SKILL.md for more
//
// server_tokens (text only, has_mtmd = false)
//
static void test_server_tokens_basic(testing & t) {
server_tokens st(llama_tokens{1, 2, 3}, false);
t.assert_equal("size", (size_t) 3, st.size());
t.assert_true("not empty", !st.empty());
t.assert_equal("operator[]", 2, st[1]);
st.push_back(4);
t.assert_equal("push_back", (size_t) 4, st.size());
t.assert_true("get_tokens", st.get_tokens() == llama_tokens({1, 2, 3, 4}));
st.insert({5, 6});
t.assert_true("insert appends", st.get_tokens() == llama_tokens({1, 2, 3, 4, 5, 6}));
st.set_token(0, 9);
t.assert_equal("set_token", 9, st[0]);
bool thrown = false;
try {
st.push_back(LLAMA_TOKEN_NULL);
} catch (const std::exception &) {
thrown = true;
}
t.assert_true("push_back(LLAMA_TOKEN_NULL) throws", thrown);
t.assert_equal("size unchanged after throw", (size_t) 6, st.size());
st.clear();
t.assert_true("clear", st.empty());
}
static void test_server_tokens_pos(testing & t) {
server_tokens st(llama_tokens{1, 2, 3, 4, 5}, false);
t.assert_equal("pos_next()", (llama_pos) 5, st.pos_next());
t.assert_equal("pos_next(0)", (llama_pos) 0, st.pos_next(0));
t.assert_equal("pos_next(3)", (llama_pos) 3, st.pos_next(3));
t.assert_equal("size_up_to_pos(0)", (size_t) 0, st.size_up_to_pos(0));
t.assert_equal("size_up_to_pos(3)", (size_t) 3, st.size_up_to_pos(3));
t.assert_equal("size_up_to_pos clamps", (size_t) 5, st.size_up_to_pos(100));
}
static void test_server_tokens_keep_first(testing & t) {
server_tokens st(llama_tokens{1, 2, 3, 4, 5}, false);
st.keep_first(5);
t.assert_equal("keep all", (size_t) 5, st.size());
st.keep_first(2);
t.assert_true("keep 2", st.get_tokens() == llama_tokens({1, 2}));
st.keep_first(0);
t.assert_true("keep 0", st.empty());
}
static void test_server_tokens_common_prefix(testing & t) {
server_tokens a(llama_tokens{1, 2, 3, 4}, false);
server_tokens empty(llama_tokens{}, false);
t.assert_equal("both empty", (size_t) 0, empty.get_common_prefix(empty));
t.assert_equal("one empty", (size_t) 0, a.get_common_prefix(empty));
t.assert_equal("identical", (size_t) 4, a.get_common_prefix(a));
t.assert_equal("diverge at 2", (size_t) 2, a.get_common_prefix(server_tokens(llama_tokens{1, 2, 9, 4}, false)));
t.assert_equal("diverge at 0", (size_t) 0, a.get_common_prefix(server_tokens(llama_tokens{9, 2, 3, 4}, false)));
t.assert_equal("b is prefix", (size_t) 2, a.get_common_prefix(server_tokens(llama_tokens{1, 2}, false)));
t.assert_equal("a is prefix", (size_t) 4, a.get_common_prefix(server_tokens(llama_tokens{1, 2, 3, 4, 5, 6}, false)));
}
static void test_server_tokens_serialize(testing & t) {
// same packing as slot save/restore in server-context.cpp
auto pack = [](const std::vector<char> & data) {
llama_tokens packed(data.size() / sizeof(llama_token));
memcpy(packed.data(), data.data(), data.size());
return packed;
};
t.test("roundtrip", [&](testing & t) {
server_tokens st(llama_tokens{10, 20, 30}, false);
std::vector<char> data = st.serialize();
t.assert_equal("aligned to llama_token", (size_t) 0, data.size() % sizeof(llama_token));
server_tokens restored = server_tokens::deserialize(pack(data), false);
t.assert_true("tokens", restored.get_tokens() == st.get_tokens());
});
t.test("roundtrip empty", [&](testing & t) {
server_tokens st(llama_tokens{}, false);
server_tokens restored = server_tokens::deserialize(pack(st.serialize()), false);
t.assert_true("empty", restored.empty());
});
t.test("legacy plain token list", [&](testing & t) {
server_tokens restored = server_tokens::deserialize(llama_tokens{7, 8, 9}, false);
t.assert_true("tokens", restored.get_tokens() == llama_tokens({7, 8, 9}));
});
auto expect_throw = [](testing & t, const std::string & msg, const llama_tokens & packed) {
bool thrown = false;
try {
server_tokens::deserialize(packed, false);
} catch (const std::exception &) {
thrown = true;
}
t.assert_true(msg, thrown);
};
t.test("corrupted", [&](testing & t) {
llama_tokens good = pack(server_tokens(llama_tokens{10, 20, 30}, false).serialize());
llama_tokens bad_version = good;
bad_version[1] = 999;
expect_throw(t, "bad version", bad_version);
llama_tokens truncated = good;
truncated.pop_back();
expect_throw(t, "truncated", truncated);
llama_tokens trailing = good;
trailing.push_back(42);
expect_throw(t, "trailing data", trailing);
llama_tokens huge_count = good;
huge_count[2] = 0x7fffffff;
expect_throw(t, "huge token count", huge_count);
expect_throw(t, "marker only", llama_tokens{LLAMA_TOKEN_NULL});
});
}
static void test_server_tokens_clone(testing & t) {
server_tokens st(llama_tokens{1, 2, 3}, false);
server_tokens cl = st.clone();
cl.push_back(4);
t.assert_equal("original untouched", (size_t) 3, st.size());
t.assert_equal("clone grew", (size_t) 4, cl.size());
t.assert_true("get_text_tokens", cl.get_text_tokens() == llama_tokens({1, 2, 3, 4}));
}
static void test_server_tokens(testing & t) {
t.test("basic", test_server_tokens_basic);
t.test("pos", test_server_tokens_pos);
t.test("keep_first", test_server_tokens_keep_first);
t.test("common_prefix", test_server_tokens_common_prefix);
t.test("serialize", test_server_tokens_serialize);
t.test("clone", test_server_tokens_clone);
}
//
// string / json helpers
//
static void test_utf8(testing & t) {
t.test("validate_utf8", [](testing & t) {
t.assert_equal("empty", (size_t) 0, validate_utf8(""));
t.assert_equal("ascii", (size_t) 3, validate_utf8("abc"));
t.assert_equal("full 2-byte", (size_t) 3, validate_utf8("a\xC3\xA9"));
t.assert_equal("full 3-byte", (size_t) 4, validate_utf8("a\xE2\x82\xAC"));
t.assert_equal("full 4-byte", (size_t) 5, validate_utf8("a\xF0\x9F\x98\x80"));
t.assert_equal("cut 2-byte (1/2)", (size_t) 1, validate_utf8("a\xC3"));
t.assert_equal("cut 3-byte (1/3)", (size_t) 1, validate_utf8("a\xE2"));
t.assert_equal("cut 3-byte (2/3)", (size_t) 1, validate_utf8("a\xE2\x82"));
t.assert_equal("cut 4-byte (1/4)", (size_t) 1, validate_utf8("a\xF0"));
t.assert_equal("cut 4-byte (2/4)", (size_t) 1, validate_utf8("a\xF0\x9F"));
t.assert_equal("cut 4-byte (3/4)", (size_t) 1, validate_utf8("a\xF0\x9F\x98"));
t.assert_equal("cut only", (size_t) 0, validate_utf8("\xF0\x9F\x98"));
t.assert_equal("full then cut", (size_t) 4, validate_utf8("\xF0\x9F\x98\x80\xC3"));
});
t.test("is_valid_utf8", [](testing & t) {
t.assert_true("empty", is_valid_utf8(""));
t.assert_true("ascii", is_valid_utf8("hello"));
t.assert_true("multi-byte", is_valid_utf8("\xC3\xA9\xE2\x82\xAC\xF0\x9F\x98\x80"));
t.assert_true("truncated", !is_valid_utf8("\xE2\x82"));
t.assert_true("bad cont. byte", !is_valid_utf8("\xC3\x41"));
t.assert_true("lone cont. byte", !is_valid_utf8("\x80"));
t.assert_true("invalid lead", !is_valid_utf8("\xFF"));
});
}
static void test_sse(testing & t) {
t.test("oai", [](testing & t) {
t.assert_equal("single", std::string("data: {\"a\":1}\n\n"), format_oai_sse(json{{"a", 1}}));
t.assert_equal("array", std::string("data: {\"a\":1}\n\ndata: {\"b\":2}\n\n"),
format_oai_sse(json::array({json{{"a", 1}}, json{{"b", 2}}})));
t.assert_equal("invalid utf8 replaced", std::string("data: \"x\xEF\xBF\xBD\"\n\n"), format_oai_sse(json("x\xC3")));
});
t.test("oai_resp", [](testing & t) {
json ev = {{"event", "response.created"}, {"data", {{"id", "r1"}}}};
t.assert_equal("single", std::string("event: response.created\ndata: {\"id\":\"r1\"}\n\n"), format_oai_resp_sse(ev));
t.assert_equal("array", std::string("event: response.created\ndata: {\"id\":\"r1\"}\n\nevent: response.created\ndata: {\"id\":\"r1\"}\n\n"),
format_oai_resp_sse(json::array({ev, ev})));
});
t.test("anthropic", [](testing & t) {
json ev = {{"event", "message_start"}, {"data", {{"type", "message_start"}}}};
t.assert_equal("with event", std::string("event: message_start\ndata: {\"type\":\"message_start\"}\n\n"), format_anthropic_sse(ev));
t.assert_equal("data only", std::string("data: {\"type\":\"ping\"}\n\n"), format_anthropic_sse(json{{"type", "ping"}}));
});
}
static void test_json_helpers(testing & t) {
t.test("json_is_array_of_numbers", [](testing & t) {
t.assert_true("ints", json_is_array_of_numbers(json::array({1, 2, 3})));
t.assert_true("empty array", json_is_array_of_numbers(json::array()));
t.assert_true("float", !json_is_array_of_numbers(json::array({1, 2.5})));
t.assert_true("string", !json_is_array_of_numbers(json::array({1, "a"})));
t.assert_true("not an array", !json_is_array_of_numbers(json(1)));
});
t.test("json_is_array_of_mixed_numbers_strings", [](testing & t) {
t.assert_true("mixed", json_is_array_of_mixed_numbers_strings(json::array({"a", 1})));
t.assert_true("numbers only", !json_is_array_of_mixed_numbers_strings(json::array({1, 2})));
t.assert_true("strings only", !json_is_array_of_mixed_numbers_strings(json::array({"a", "b"})));
t.assert_true("empty", !json_is_array_of_mixed_numbers_strings(json::array()));
});
t.test("json_is_array_and_contains_numbers", [](testing & t) {
t.assert_true("contains", json_is_array_and_contains_numbers(json::array({"a", 1})));
t.assert_true("strings only", !json_is_array_and_contains_numbers(json::array({"a"})));
t.assert_true("not an array", !json_is_array_and_contains_numbers(json(1)));
});
t.test("json_get_nested_values", [](testing & t) {
json js = {{"a", {{"b", {{"c", 1}}}}}, {"x", 2}};
json res = json_get_nested_values({"a/b/c", "x", "a/b", "a/missing", "x/y"}, js);
t.assert_equal("nested", 1, res.at("a/b/c").get<int>());
t.assert_equal("top", 2, res.at("x").get<int>());
t.assert_true("object", res.at("a/b") == json{{"c", 1}});
t.assert_true("missing", !res.contains("a/missing"));
t.assert_true("not obj", !res.contains("x/y"));
t.assert_equal("count", (size_t) 3, res.size());
});
}
static void test_lora_helpers(testing & t) {
// only pointer identity is used, never dereferenced
auto * p0 = reinterpret_cast<llama_adapter_lora *>(0x10);
auto * p1 = reinterpret_cast<llama_adapter_lora *>(0x20);
auto make = [](llama_adapter_lora * ptr, float scale) {
common_adapter_lora_info info;
info.scale = scale;
info.ptr = ptr;
return info;
};
t.test("parse_lora_request", [](testing & t) {
auto res = parse_lora_request(json::array({
json{{"id", 0}, {"scale", 0.5}},
json{{"id", 2}, {"scale", 1.0}},
json{{"id", 0}, {"scale", 0.25}}, // last one wins
}));
t.assert_equal("count", (size_t) 2, res.size());
t.assert_equal("id 0", 0.25f, res.at(0));
t.assert_equal("id 2", 1.0f, res.at(2));
auto defaults = parse_lora_request(json::array({json::object()}));
t.assert_equal("default id/scale", 0.0f, defaults.at(-1));
});
t.test("are_lora_equal", [&](testing & t) {
t.assert_true("both empty", are_lora_equal({}, {}));
t.assert_true("same", are_lora_equal({make(p0, 1.0f)}, {make(p0, 1.0f)}));
t.assert_true("diff scale", !are_lora_equal({make(p0, 1.0f)}, {make(p0, 0.5f)}));
t.assert_true("diff ptr", !are_lora_equal({make(p0, 1.0f)}, {make(p1, 1.0f)}));
t.assert_true("diff size", !are_lora_equal({make(p0, 1.0f)}, {}));
});
t.test("lora_get_enabled_ids", [&](testing & t) {
auto ids = lora_get_enabled_ids({make(p0, 0.0f), make(p1, 0.5f), make(p0, -1.0f), make(p1, 2.0f)});
t.assert_true("ids", ids == std::vector<size_t>({1, 3}));
});
}
//
// server_pipe
//
static void test_server_pipe(testing & t) {
auto stop_now = []() { return true; };
t.test("drain then eof", [](testing & t) {
server_pipe<int> p;
t.assert_true("write 1", p.write(1));
t.assert_true("write 2", p.write(2));
p.close_write();
int v = 0;
t.assert_true("read 1", p.read(v, nullptr) && v == 1);
t.assert_true("read 2", p.read(v, nullptr) && v == 2);
t.assert_true("eof", !p.read(v, nullptr));
});
t.test("broken pipe", [](testing & t) {
server_pipe<int> p;
p.close_read();
t.assert_true("write fails", !p.write(1));
});
t.test("max_size drops oldest", [](testing & t) {
server_pipe<int> p;
p.max_size = 2;
p.write(1);
p.write(2);
p.write(3);
p.close_write();
int v = 0;
t.assert_true("read 2", p.read(v, nullptr) && v == 2);
t.assert_true("read 3", p.read(v, nullptr) && v == 3);
t.assert_true("eof", !p.read(v, nullptr));
});
t.test("queued data wins over should_stop", [&](testing & t) {
server_pipe<int> p;
p.write(1);
int v = 0;
t.assert_true("read", p.read(v, stop_now) && v == 1);
});
t.test("should_stop closes the pipe", [&](testing & t) {
server_pipe<int> p;
int v = 0;
t.assert_true("read stopped", !p.read(v, stop_now));
t.assert_true("write fails", !p.write(1));
});
t.test("should_stop as deadline keeps the pipe open", [&](testing & t) {
server_pipe<int> p;
int v = 0;
t.assert_true("read stopped", !p.read(v, stop_now, /* close_on_stop */ false));
t.assert_true("write ok", p.write(1));
t.assert_true("read again", p.read(v, nullptr) && v == 1);
});
t.test("cross thread", [](testing & t) {
constexpr int n = 1000;
server_pipe<int> p;
std::thread producer([&]() {
for (int i = 0; i < n; i++) {
p.write(int(i));
}
p.close_write();
});
int v = 0;
int count = 0;
bool in_order = true;
while (p.read(v, nullptr)) {
in_order &= v == count;
count++;
}
producer.join();
t.assert_equal("count", n, count);
t.assert_true("in order", in_order);
});
}
//
// server_decision_escape_special_tokens
//
// "<|name|>" with U+00A6 in place of "|"
static std::string escaped(const std::string & name) {
return "<\xC2\xA6" + name + "\xC2\xA6>";
}
static void test_decision_escape(testing & t) {
t.test("cases", [](testing & t) {
const std::vector<std::pair<std::string, std::string>> cases = {
{ "", "" },
{ "no tokens", "no tokens" },
{ "<|im_start|>user", escaped("im_start") + "user" },
{ "<|a_1|><|B|>", escaped("a_1") + escaped("B") },
{ "<|<|a|>", "<|" + escaped("a") },
{ "<|a|>|>", escaped("a") + "|>" },
{ "<||>", "<||>" },
{ "<|a b|>", "<|a b|>" },
{ "<|a|b|>", "<|a|b|>" },
{ "<|caf\xC3\xA9|>", "<|caf\xC3\xA9|>" },
{ "<|abc", "<|abc" },
};
for (const auto & [text, expected] : cases) {
t.assert_equal(text, expected, server_decision_escape_special_tokens(text));
}
});
// every string of up to 6 chars from "<|>a-" gives the text of the regex
t.test("same as std::regex", [](testing & t) {
static const std::regex re_special("<\\|([A-Za-z0-9_]+)\\|>");
const std::string chars = "<|>a-";
size_t n_strings = 0;
size_t n_diff = 0;
for (size_t len = 0; len <= 6; len++) {
size_t n = 1;
for (size_t k = 0; k < len; k++) {
n *= chars.size();
}
for (size_t idx = 0; idx < n; idx++) {
std::string text;
for (size_t k = 0, v = idx; k < len; k++, v /= chars.size()) {
text += chars[v % chars.size()];
}
if (server_decision_escape_special_tokens(text) != std::regex_replace(text, re_special, "<\xC2\xA6$1\xC2\xA6>")) {
n_diff++;
}
n_strings++;
}
}
t.assert_equal("strings that differ, of " + std::to_string(n_strings), (size_t) 0, n_diff);
});
// std::regex overflowed the stack on these
t.test("long input", [](testing & t) {
const std::string name(1000000, 'a');
t.assert_true("not closed", server_decision_escape_special_tokens("<|" + name) == "<|" + name);
t.assert_true("closed", server_decision_escape_special_tokens("<|" + name + "|>") == escaped(name));
});
}
int main(int argc, char ** argv) {
testing t;
const char * verbose = getenv("LLAMA_TEST_VERBOSE");
if (verbose) {
t.verbose = std::string(verbose) == "1";
}
if (argc > 1) {
t.set_filter(argv[1]);
}
t.test("server_tokens", test_server_tokens);
t.test("utf8", test_utf8);
t.test("sse", test_sse);
t.test("json_helpers", test_json_helpers);
t.test("lora_helpers", test_lora_helpers);
t.test("server_pipe", test_server_pipe);
t.test("decision_escape", test_decision_escape);
return t.summary();
}