mirror of
https://github.com/LostRuins/koboldcpp.git
synced 2026-10-21 22:00:02 +00:00
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-".
459 lines
18 KiB
C++
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();
|
|
}
|