From d67c0b4107112e4790774c3a8169e2e3eb24643b Mon Sep 17 00:00:00 2001 From: helanfxz <126638465+helanfxz@users.noreply.github.com> Date: Sat, 25 Jul 2026 16:23:31 +0800 Subject: [PATCH 1/9] tests: synchronize save-load-state generation (#26056) --- tests/test-save-load-state.cpp | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test-save-load-state.cpp b/tests/test-save-load-state.cpp index 6e93ce6fb8..bbb025617f 100644 --- a/tests/test-save-load-state.cpp +++ b/tests/test-save-load-state.cpp @@ -44,6 +44,8 @@ static llama_tokens generate_tokens(llama_context * ctx, llama_sampler * smpl, i n_past++; } + llama_synchronize(ctx); + return result; } From 910196f6b3dfc6aca88fa732e2b02f270ff9b56b Mon Sep 17 00:00:00 2001 From: Aldehir Rojas Date: Sat, 25 Jul 2026 04:58:09 -0500 Subject: [PATCH 2/9] common : add support for multiple end sequences in the reasoning budget sampler (#25544) * common : extract trie/ac to a separate file * common : support multiple token sequences in the reasoning budget sampler * common/trie : return matched word index * common/trie : rename "word" to "pattern" * common/reasoning-budget : expose matched end sequence * common/sampling : replay end sequence when reasoning budget is done * cont : update to use multiple end sequences * cont : clean up --- common/CMakeLists.txt | 2 + common/chat.cpp | 22 +++-- common/chat.h | 2 +- common/common.h | 12 +-- common/peg-parser.cpp | 158 +------------------------------- common/reasoning-budget.cpp | 116 +++++++++++++++-------- common/reasoning-budget.h | 26 ++++-- common/sampling.cpp | 13 ++- common/trie.cpp | 123 +++++++++++++++++++++++++ common/trie.h | 73 +++++++++++++++ tests/test-chat.cpp | 16 ++-- tests/test-reasoning-budget.cpp | 149 ++++++++++++++++++++++++++---- tools/server/server-common.cpp | 4 +- tools/server/server-schema.cpp | 35 +++++-- 14 files changed, 497 insertions(+), 254 deletions(-) create mode 100644 common/trie.cpp create mode 100644 common/trie.h diff --git a/common/CMakeLists.txt b/common/CMakeLists.txt index 4cf580a056..99688f53b8 100644 --- a/common/CMakeLists.txt +++ b/common/CMakeLists.txt @@ -100,6 +100,8 @@ add_library(${TARGET} sampling.h speculative.cpp speculative.h + trie.cpp + trie.h unicode.cpp unicode.h jinja/lexer.cpp diff --git a/common/chat.cpp b/common/chat.cpp index 2b0b40ccbe..7a6e7238cf 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -1024,7 +1024,7 @@ static common_chat_params common_chat_params_init_ministral_3(const common_chat_ data.supports_thinking = true; data.thinking_start_tag = "[THINK]"; - data.thinking_end_tag = "[/THINK]"; + data.thinking_end_tags = {"[/THINK]"}; data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, /* messages_override = */ adjusted_messages); data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs, /* messages_override = */ adjusted_messages); data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; @@ -1150,6 +1150,9 @@ static common_chat_params common_chat_params_init_gpt_oss(const common_chat_temp data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.supports_thinking = true; + data.thinking_start_tag = "<|channel|>analysis<|message|>"; + data.thinking_end_tags = {"<|end|>"}; + // These special tokens are required to parse properly, so we include them // even if parse_tool_calls is false. data.preserved_tokens = { @@ -1294,7 +1297,7 @@ static common_chat_params common_chat_params_init_gemma4(const common_chat_templ data.format = COMMON_CHAT_FORMAT_PEG_GEMMA4; data.supports_thinking = true; data.thinking_start_tag = "<|channel>thought"; - data.thinking_end_tag = ""; + data.thinking_end_tags = {""}; data.preserved_tokens = { "<|channel>", @@ -1569,7 +1572,7 @@ static common_chat_params common_chat_params_init_kimi_k2(const common_chat_temp const std::string GEN_PROMPT = "<|im_assistant|>assistant<|im_middle|>"; data.thinking_start_tag = THINK_START; - data.thinking_end_tag = THINK_END; + data.thinking_end_tags = {THINK_END}; if (inputs.has_continuation()) { const auto & msg = inputs.continue_msg; @@ -1703,7 +1706,7 @@ static common_chat_params common_chat_params_init_lfm2(const common_chat_templat } data.thinking_start_tag = THINK_START; - data.thinking_end_tag = THINK_END; + data.thinking_end_tags = {THINK_END}; auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); auto has_response_format = !inputs.json_schema.is_null() && inputs.json_schema.is_object(); @@ -1943,7 +1946,7 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.supports_thinking = true; data.thinking_start_tag = ""; - data.thinking_end_tag = ""; + data.thinking_end_tags = {""}; data.preserved_tokens = { "｜DSML｜", "", @@ -2160,7 +2163,7 @@ static common_chat_params common_chat_params_init_cohere2moe(const common_chat_t data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.supports_thinking = true; data.thinking_start_tag = THINK_START; - data.thinking_end_tag = THINK_END; + data.thinking_end_tags = {THINK_END}; data.preserved_tokens = { TURN_START, TURN_END, CHATBOT, USER, SYSTEM, THINK_START, THINK_END, @@ -2510,7 +2513,7 @@ static common_chat_params common_chat_params_init_minicpm5(const common_chat_tem }; data.thinking_start_tag = ""; - data.thinking_end_tag = ""; + data.thinking_end_tags = {""}; data.message_delimiters = { { COMMON_CHAT_ROLE_ASSISTANT, "<|im_start|>assistant" }, @@ -2866,7 +2869,10 @@ static common_chat_params common_chat_templates_apply_jinja(const struct common_ auto_params.supports_thinking = autoparser.reasoning.mode != autoparser::reasoning_mode::NONE; if (auto_params.supports_thinking) { auto_params.thinking_start_tag = trim_whitespace(autoparser.reasoning.start); - auto_params.thinking_end_tag = trim_whitespace(autoparser.reasoning.end); + auto end_tag = trim_whitespace(autoparser.reasoning.end); + if (!end_tag.empty()) { + auto_params.thinking_end_tags = {std::move(end_tag)}; + } } common_peg_arena arena; arena.load(auto_params.parser); diff --git a/common/chat.h b/common/chat.h index 7898f1623f..d79f4ecd77 100644 --- a/common/chat.h +++ b/common/chat.h @@ -274,7 +274,7 @@ struct common_chat_params { std::string generation_prompt; bool supports_thinking = false; std::string thinking_start_tag; // e.g., "" - std::string thinking_end_tag; // e.g., "" + std::vector thinking_end_tags; // e.g., "" std::vector grammar_triggers; std::vector preserved_tokens; std::vector additional_stops; diff --git a/common/common.h b/common/common.h index cb1827f5dd..b5687c1083 100644 --- a/common/common.h +++ b/common/common.h @@ -284,12 +284,12 @@ struct common_params_sampling { // reasoning budget sampler parameters // these are populated by the server/CLI based on chat template params - int32_t reasoning_budget_tokens = -1; // -1 = disabled, >= 0 = token budget - std::vector reasoning_budget_start; // start tag token sequence - std::vector reasoning_budget_end; // end tag token sequence - std::vector reasoning_budget_forced; // forced sequence (message + end tag) - std::string reasoning_budget_message; // message injected before end tag when budget exhausted - bool reasoning_control = false; // create the budget sampler on demand so reasoning can be ended at runtime + int32_t reasoning_budget_tokens = -1; // -1 = disabled, >= 0 = token budget + std::vector reasoning_budget_start; // start tag token sequence + std::vector reasoning_budget_end; // end tag token sequences; the first tag is used as the forcing sequence + std::vector reasoning_budget_forced; // forced sequence (message + first end tag) + std::string reasoning_budget_message; // message injected before end tag when budget exhausted + bool reasoning_control = false; // create the budget sampler on demand so reasoning can be ended at runtime bool backend_sampling = false; diff --git a/common/peg-parser.cpp b/common/peg-parser.cpp index 807e952d90..ef290ed7c0 100644 --- a/common/peg-parser.cpp +++ b/common/peg-parser.cpp @@ -3,10 +3,10 @@ #include "common.h" #include "json-schema-to-grammar.h" #include "log.h" +#include "trie.h" #include "unicode.h" #include -#include #include #include #include @@ -32,154 +32,6 @@ static bool is_hex_digit(const char c) { return (c >= '0' && c <= '9') || (c >= 'a' && c <= 'f') || (c >= 'A' && c <= 'F'); } -// Trie for matching multiple literals. -// This is used in common_peg_until_parser and to build a GBNF exclusion grammar -struct trie { - struct node { - std::map children; // Use uint32_t to store Unicode codepoints - bool is_word; - }; - - std::vector nodes; - - trie(const std::vector & words) { - create_node(); // root node - for (const auto & w : words) { - insert(w); - } - } - - enum match_result { NO_MATCH, PARTIAL_MATCH, COMPLETE_MATCH }; - - // Check if a delimiter starts at the given position - match_result check_at(std::string_view sv, size_t start_pos) const { - size_t current = 0; // Start at root - size_t pos = start_pos; - - // LOG_DBG("%s: checking at pos %zu, sv='%s'\n", __func__, start_pos, std::string(sv).c_str()); - - while (pos < sv.size()) { - auto result = common_parse_utf8_codepoint(sv, pos); - if (result.status != utf8_parse_result::SUCCESS) { - break; - } - - auto it = nodes[current].children.find(result.codepoint); - if (it == nodes[current].children.end()) { - // Can't continue matching - return match_result{match_result::NO_MATCH}; - } - - current = it->second; - pos += result.bytes_consumed; - - // Check if we've matched a complete word - if (nodes[current].is_word) { - return match_result{match_result::COMPLETE_MATCH}; - } - } - - // Reached end of input while still in the trie (not at root) - if (current != 0) { - // We're in the middle of a potential match - return match_result{match_result::PARTIAL_MATCH}; - } - - // Reached end at root (no match) - return match_result{match_result::NO_MATCH}; - } - - private: - size_t create_node() { - size_t index = nodes.size(); - nodes.emplace_back(); - return index; - } - - void insert(const std::string & word) { - size_t current = 0; - size_t pos = 0; - while (pos < word.length()) { - auto result = common_parse_utf8_codepoint(word, pos); - if (result.status != utf8_parse_result::SUCCESS) { - break; - } - - uint32_t ch = result.codepoint; - pos += result.bytes_consumed; - - auto it = nodes[current].children.find(ch); - if (it == nodes[current].children.end()) { - size_t child = create_node(); - nodes[current].children[ch] = child; - current = child; - } else { - current = it->second; - } - } - nodes[current].is_word = true; - } -}; - -// Aho-Corasick automaton -struct aho_corasick { - trie t; - std::vector fail; // failure links - std::vector order; // states in BFS order - std::vector terminal; // match states (directly or via a suffix link) - std::set alphabet; // every character with a transition - - aho_corasick(const std::vector & strings) : t(strings) { - const auto & nodes = t.nodes; - const size_t n = nodes.size(); - - fail.assign(n, 0); - order.reserve(n); - - std::deque queue{ 0 }; - while (!queue.empty()) { - size_t u = queue.front(); - queue.pop_front(); - order.push_back(u); - for (const auto & [ch, v] : nodes[u].children) { - if (u != 0) { - size_t f = fail[u]; - while (f && nodes[f].children.find(ch) == nodes[f].children.end()) { - f = fail[f]; - } - auto it = nodes[f].children.find(ch); - fail[v] = (it != nodes[f].children.end() && it->second != v) ? it->second : 0; - } - queue.push_back(v); - } - } - - terminal.assign(n, false); - for (size_t u : order) { - terminal[u] = nodes[u].is_word || (u != 0 && terminal[fail[u]]); - } - - for (const auto & node : nodes) { - for (const auto & [ch, v] : node.children) { - alphabet.insert(ch); - } - } - } - - size_t num_states() const { return t.nodes.size(); } - bool is_terminal(size_t s) const { return terminal[s]; } - - // follow failure links until a transition on `ch` exists. - size_t next(size_t state, uint32_t ch) const { - const auto & nodes = t.nodes; - while (state && nodes[state].children.find(ch) == nodes[state].children.end()) { - state = fail[state]; - } - auto it = nodes[state].children.find(ch); - return it != nodes[state].children.end() ? it->second : 0; - } -}; - static std::pair parse_hex_escape(const std::string & str, size_t pos, int hex_count) { if (pos + hex_count > str.length()) { return {0, 0}; @@ -797,7 +649,7 @@ struct parser_executor { } common_peg_parse_result operator()(const common_peg_until_parser & p) const { - trie matcher(p.delimiters); + common_trie matcher(p.delimiters); // Scan input and check for delimiters size_t pos = start_pos; @@ -824,12 +676,12 @@ struct parser_executor { // Check if a delimiter starts at this position auto match = matcher.check_at(ctx.input, pos); - if (match == trie::COMPLETE_MATCH) { + if (match == common_trie::COMPLETE_MATCH) { // Found a complete delimiter, return everything before it return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos); } - if (match == trie::PARTIAL_MATCH) { + if (match == common_trie::PARTIAL_MATCH) { // Found a partial match extending to end of input, return everything before it return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos); } @@ -1559,7 +1411,7 @@ static std::string gbnf_ac_grammar( const std::map> &, const std::vector &, const std::function &)> & build_rule) { - aho_corasick ac(strings); + common_aho_corasick ac(strings); auto state_name = [&](size_t s) -> std::string { if (s == 0) { diff --git a/common/reasoning-budget.cpp b/common/reasoning-budget.cpp index 7da0bb1c57..1fe242d062 100644 --- a/common/reasoning-budget.cpp +++ b/common/reasoning-budget.cpp @@ -1,39 +1,52 @@ #include "reasoning-budget.h" #include "common.h" +#include "trie.h" #include "unicode.h" #include "log.h" +#include #include #include #include #include struct token_matcher { - std::vector tokens; - size_t pos = 0; + std::vector seqs; + common_aho_corasick ac; + size_t state = 0; - bool advance(llama_token token) { - if (tokens.empty()) { - return false; - } + token_matcher(const std::vector & seqs) : seqs(collect(seqs)), ac(build_trie(this->seqs)) {} - if (token == tokens[pos]) { - pos++; - if (pos >= tokens.size()) { - pos = 0; - return true; - } - } else { - pos = 0; - if (token == tokens[0]) { - pos = 1; + static std::vector collect(const std::vector & seqs) { + std::vector res; + for (const auto & seq : seqs) { + if (!seq.empty() && std::find(res.begin(), res.end(), seq) == res.end()) { + res.push_back(seq); } } - return false; + return res; } - void reset() { pos = 0; } + static common_trie build_trie(const std::vector & seqs) { + common_trie t; + for (const auto & seq : seqs) { + t.insert(std::vector(seq.begin(), seq.end())); + } + return t; + } + + // returns the index into seqs of the longest sequence ending at this token, or -1 + int32_t advance(llama_token token) { + state = ac.next(state, (uint32_t) token); + const int32_t p = ac.match_pattern(state); + if (p >= 0) { + state = 0; + } + return p; + } + + void reset() { state = 0; } }; struct common_reasoning_budget_ctx { @@ -41,7 +54,7 @@ struct common_reasoning_budget_ctx { token_matcher start_matcher; token_matcher end_matcher; - std::vector forced_tokens; + llama_tokens forced_tokens; int32_t budget; // maximum tokens in reasoning block int32_t remaining; // tokens remaining in budget @@ -50,6 +63,8 @@ struct common_reasoning_budget_ctx { // for forcing size_t force_pos; // next position in forced_tokens to force + + int32_t end_match; // index into end_matcher.seqs of the sequence that transitioned to DONE, -1 if none }; static const char * common_reasoning_budget_name(const struct llama_sampler * /*smpl*/) { @@ -62,7 +77,7 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to switch (ctx->state) { case REASONING_BUDGET_IDLE: { - if (ctx->start_matcher.advance(token)) { + if (ctx->start_matcher.advance(token) >= 0) { ctx->state = REASONING_BUDGET_COUNTING; ctx->remaining = ctx->budget; COM_TRC("activated, budget=%d tokens\n", ctx->budget); @@ -78,8 +93,10 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to case REASONING_BUDGET_COUNTING: case REASONING_BUDGET_WAITING_UTF8: { - if (ctx->end_matcher.advance(token)) { + const int32_t match = ctx->end_matcher.advance(token); + if (match >= 0) { ctx->state = REASONING_BUDGET_DONE; + ctx->end_match = match; COM_TRC("%s", "deactivated (natural end)\n"); break; } @@ -115,19 +132,25 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to break; } case REASONING_BUDGET_FORCING: + { + // track the end sequence within forced_tokens so it is also reported on DONE + const int32_t match = ctx->end_matcher.advance(token); ctx->force_pos++; if (ctx->force_pos >= ctx->forced_tokens.size()) { ctx->state = REASONING_BUDGET_DONE; + ctx->end_match = match; COM_TRC("%s", "forced sequence complete, done\n"); } break; + } case REASONING_BUDGET_DONE: // Re-arm on a new start tag: some models emit multiple blocks // per response, and each should get a fresh budget window. - if (ctx->start_matcher.advance(token)) { + if (ctx->start_matcher.advance(token) >= 0) { ctx->state = REASONING_BUDGET_COUNTING; ctx->remaining = ctx->budget; ctx->end_matcher.reset(); + ctx->end_match = -1; COM_TRC("re-activated on new start tag, budget=%d tokens\n", ctx->budget); if (ctx->remaining <= 0) { @@ -169,11 +192,12 @@ static void common_reasoning_budget_reset(struct llama_sampler * smpl) { ctx->start_matcher.reset(); ctx->end_matcher.reset(); ctx->force_pos = 0; + ctx->end_match = -1; } static struct llama_sampler * common_reasoning_budget_init_state( - const struct llama_vocab * vocab, const std::vector & start_tokens, - const std::vector & end_tokens, const std::vector & forced_tokens, + const struct llama_vocab * vocab, const std::vector & start_seqs, + const std::vector & end_seqs, const llama_tokens & forced_tokens, int32_t budget, common_reasoning_budget_state initial_state); static struct llama_sampler * common_reasoning_budget_clone(const struct llama_sampler * smpl); @@ -205,12 +229,12 @@ static struct llama_sampler * common_reasoning_budget_clone(const struct llama_s } static struct llama_sampler * common_reasoning_budget_init_state( - const struct llama_vocab * vocab, - const std::vector & start_tokens, - const std::vector & end_tokens, - const std::vector & forced_tokens, - int32_t budget, - common_reasoning_budget_state initial_state) { + const struct llama_vocab * vocab, + const std::vector & start_seqs, + const std::vector & end_seqs, + const llama_tokens & forced_tokens, + int32_t budget, + common_reasoning_budget_state initial_state) { // promote COUNTING with budget <= 0 to FORCING if (initial_state == REASONING_BUDGET_COUNTING && budget <= 0) { initial_state = REASONING_BUDGET_FORCING; @@ -220,25 +244,26 @@ static struct llama_sampler * common_reasoning_budget_init_state( /* .iface = */ &common_reasoning_budget_i, /* .ctx = */ new common_reasoning_budget_ctx { /* .vocab = */ vocab, - /* .start_matcher = */ { start_tokens, 0 }, - /* .end_matcher = */ { end_tokens, 0 }, + /* .start_matcher = */ token_matcher(start_seqs), + /* .end_matcher = */ token_matcher(end_seqs), /* .forced_tokens = */ forced_tokens, /* .budget = */ budget, /* .remaining = */ budget, /* .state = */ initial_state, /* .force_pos = */ 0, + /* .end_match = */ -1, } ); } struct llama_sampler * common_reasoning_budget_init( - const struct llama_vocab * vocab, - const std::vector & start_tokens, - const std::vector & end_tokens, - const std::vector & forced_tokens, - int32_t budget, - common_reasoning_budget_state initial_state) { - return common_reasoning_budget_init_state(vocab, start_tokens, end_tokens, forced_tokens, budget, initial_state); + const struct llama_vocab * vocab, + const std::vector & start_seqs, + const std::vector & end_seqs, + const llama_tokens & forced_tokens, + int32_t budget, + common_reasoning_budget_state initial_state) { + return common_reasoning_budget_init_state(vocab, start_seqs, end_seqs, forced_tokens, budget, initial_state); } common_reasoning_budget_state common_reasoning_budget_get_state(const struct llama_sampler * smpl) { @@ -248,6 +273,19 @@ common_reasoning_budget_state common_reasoning_budget_get_state(const struct lla return ((const common_reasoning_budget_ctx *)smpl->ctx)->state; } +const llama_tokens * common_reasoning_budget_get_end_match(const struct llama_sampler * smpl) { + if (!smpl) { + return nullptr; + } + + const auto * ctx = (const common_reasoning_budget_ctx *) smpl->ctx; + if (ctx->end_match < 0) { + return nullptr; + } + + return &ctx->end_matcher.seqs[ctx->end_match]; +} + bool common_reasoning_budget_force(struct llama_sampler * smpl) { if (!smpl) { return false; diff --git a/common/reasoning-budget.h b/common/reasoning-budget.h index 0cf689a566..1b89a04c42 100644 --- a/common/reasoning-budget.h +++ b/common/reasoning-budget.h @@ -2,6 +2,8 @@ #include "llama.h" +#include "common.h" + #include #include @@ -17,30 +19,34 @@ enum common_reasoning_budget_state { // reasoning block (e.g. between and ). // // State machine: IDLE -> COUNTING -> WAITING_UTF8 -> FORCING -> DONE -// IDLE: passthrough, watching for start_tokens sequence -// COUNTING: counting down remaining tokens, watching for natural end_tokens +// IDLE: passthrough, watching for a start sequence +// COUNTING: counting down remaining tokens, watching for a natural end sequence // WAITING_UTF8: budget exhausted, allowing tokens to complete a UTF-8 sequence // FORCING: forces forced_tokens token-by-token (all other logits -> -inf) // DONE: passthrough forever // // Parameters: // vocab - vocabulary (used for UTF-8 boundary detection; can be nullptr) -// start_tokens - token sequence that activates counting -// end_tokens - token sequence for natural deactivation +// start_seqs - token sequences, any of which activates counting +// end_seqs - token sequences, any of which naturally deactivates // forced_tokens - token sequence forced when budget expires // budget - max tokens allowed in the reasoning block // initial_state - initial state // struct llama_sampler * common_reasoning_budget_init( - const struct llama_vocab * vocab, - const std::vector & start_tokens, - const std::vector & end_tokens, - const std::vector & forced_tokens, - int32_t budget, - common_reasoning_budget_state initial_state = REASONING_BUDGET_IDLE); + const struct llama_vocab * vocab, + const std::vector & start_seqs, + const std::vector & end_seqs, + const llama_tokens & forced_tokens, + int32_t budget, + common_reasoning_budget_state initial_state = REASONING_BUDGET_IDLE); common_reasoning_budget_state common_reasoning_budget_get_state(const struct llama_sampler * smpl); +// The end sequence that transitioned the sampler to DONE, or nullptr if none +// was recorded. Cleared when a new start sequence re-arms the sampler. +const llama_tokens * common_reasoning_budget_get_end_match(const struct llama_sampler * smpl); + // Manually transition the reasoning budget sampler into the FORCING state. // Returns true if the transition occurred. bool common_reasoning_budget_force(struct llama_sampler * smpl); diff --git a/common/sampling.cpp b/common/sampling.cpp index 75a299e23e..7b241e34f7 100644 --- a/common/sampling.cpp +++ b/common/sampling.cpp @@ -299,7 +299,7 @@ struct common_sampler * common_sampler_init(const struct llama_model * model, st if (!params.reasoning_budget_start.empty() && !params.reasoning_budget_end.empty() && (params.grammar_lazy || params.reasoning_budget_tokens >= 0 || params.reasoning_control)) { rbudget = common_reasoning_budget_init( vocab, - params.reasoning_budget_start, + {params.reasoning_budget_start}, params.reasoning_budget_end, params.reasoning_budget_forced, params.reasoning_budget_tokens < 0 ? INT_MAX : params.reasoning_budget_tokens); @@ -453,6 +453,17 @@ void common_sampler_accept(struct common_sampler * gsmpl, llama_token token, boo if (gsmpl->rbudget && is_generated) { llama_sampler_accept(gsmpl->rbudget, token); + + // if done, replay end sequence which may contain a grammar trigger + const bool is_done = common_reasoning_budget_get_state(gsmpl->rbudget) == REASONING_BUDGET_DONE; + if (gsmpl->grmr && !accept_grammar && is_done) { + const llama_tokens * end_seq = common_reasoning_budget_get_end_match(gsmpl->rbudget); + if (end_seq) { + for (const llama_token end_token : *end_seq) { + llama_sampler_accept(gsmpl->grmr, end_token); + } + } + } } if (gsmpl->grmr && accept_grammar) { diff --git a/common/trie.cpp b/common/trie.cpp new file mode 100644 index 0000000000..b5c9666ba2 --- /dev/null +++ b/common/trie.cpp @@ -0,0 +1,123 @@ +#include "trie.h" + +#include "unicode.h" + +#include + +common_trie::match_result common_trie::check_at(std::string_view sv, size_t start_pos) const { + size_t current = 0; // Start at root + size_t pos = start_pos; + + // LOG_DBG("%s: checking at pos %zu, sv='%s'\n", __func__, start_pos, std::string(sv).c_str()); + + while (pos < sv.size()) { + auto result = common_parse_utf8_codepoint(sv, pos); + if (result.status != utf8_parse_result::SUCCESS) { + break; + } + + auto it = nodes[current].children.find(result.codepoint); + if (it == nodes[current].children.end()) { + // Can't continue matching + return match_result{match_result::NO_MATCH}; + } + + current = it->second; + pos += result.bytes_consumed; + + // Check if we've matched a complete word + if (nodes[current].pattern >= 0) { + return match_result{match_result::COMPLETE_MATCH}; + } + } + + // Reached end of input while still in the trie (not at root) + if (current != 0) { + // We're in the middle of a potential match + return match_result{match_result::PARTIAL_MATCH}; + } + + // Reached end at root (no match) + return match_result{match_result::NO_MATCH}; +} + +int32_t common_trie::insert(const std::string & word) { + std::vector symbols; + size_t pos = 0; + while (pos < word.length()) { + auto result = common_parse_utf8_codepoint(word, pos); + if (result.status != utf8_parse_result::SUCCESS) { + break; + } + + symbols.push_back(result.codepoint); + pos += result.bytes_consumed; + } + return insert(symbols); +} + +int32_t common_trie::insert(const std::vector & symbols) { + size_t current = 0; + for (uint32_t ch : symbols) { + auto it = nodes[current].children.find(ch); + if (it == nodes[current].children.end()) { + size_t child = create_node(); + nodes[current].children[ch] = child; + current = child; + } else { + current = it->second; + } + } + if (nodes[current].pattern < 0) { + nodes[current].pattern = n_patterns++; + } + return nodes[current].pattern; +} + +common_aho_corasick::common_aho_corasick(common_trie trie) : t(std::move(trie)) { + const auto & nodes = t.nodes; + const size_t n = nodes.size(); + + fail.assign(n, 0); + order.reserve(n); + + std::deque queue{ 0 }; + while (!queue.empty()) { + size_t u = queue.front(); + queue.pop_front(); + order.push_back(u); + for (const auto & [ch, v] : nodes[u].children) { + if (u != 0) { + size_t f = fail[u]; + while (f && nodes[f].children.find(ch) == nodes[f].children.end()) { + f = fail[f]; + } + auto it = nodes[f].children.find(ch); + fail[v] = (it != nodes[f].children.end() && it->second != v) ? it->second : 0; + } + queue.push_back(v); + } + } + + // fail[u] points to a strictly shorter suffix, so the first pattern found on + // the fail chain (including u itself) is the longest pattern ending at u + match.assign(n, -1); + for (size_t u : order) { + match[u] = nodes[u].pattern >= 0 ? nodes[u].pattern : (u != 0 ? match[fail[u]] : -1); + } + + for (const auto & node : nodes) { + for (const auto & [ch, v] : node.children) { + alphabet.insert(ch); + } + } +} + +size_t common_aho_corasick::next(size_t state, uint32_t ch) const { + const auto & nodes = t.nodes; + while (state && nodes[state].children.find(ch) == nodes[state].children.end()) { + state = fail[state]; + } + auto it = nodes[state].children.find(ch); + return it != nodes[state].children.end() ? it->second : 0; +} diff --git a/common/trie.h b/common/trie.h new file mode 100644 index 0000000000..0f7b16a36a --- /dev/null +++ b/common/trie.h @@ -0,0 +1,73 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +// Trie for matching multiple literals. +// This is used in common_peg_until_parser and to build a GBNF exclusion grammar +struct common_trie { + struct node { + std::map children; // Use uint32_t to store Unicode codepoints + int32_t pattern = -1; // index of the pattern ending at this node, -1 if none + }; + + std::vector nodes; + + common_trie() { + create_node(); // root node + } + + common_trie(const std::vector & words) : common_trie() { + for (const auto & w : words) { + insert(w); + } + } + + enum match_result { NO_MATCH, PARTIAL_MATCH, COMPLETE_MATCH }; + + // Check if a delimiter starts at the given position + match_result check_at(std::string_view sv, size_t start_pos) const; + + // Insert a word as a sequence of Unicode codepoints, returns its pattern index + int32_t insert(const std::string & word); + + // Insert a raw symbol sequence, returns its pattern index (insertion order, + // duplicates keep the first index) + int32_t insert(const std::vector & symbols); + + private: + int32_t n_patterns = 0; + + size_t create_node() { + size_t index = nodes.size(); + nodes.emplace_back(); + return index; + } +}; + +// Aho-Corasick automaton +struct common_aho_corasick { + common_trie t; + std::vector fail; // failure links + std::vector order; // states in BFS order + std::vector match; // longest pattern ending at each state (directly or via a suffix link), -1 if none + std::set alphabet; // every character with a transition + + common_aho_corasick(common_trie trie); + + common_aho_corasick(const std::vector & strings) + : common_aho_corasick(common_trie(strings)) {} + + size_t num_states() const { return t.nodes.size(); } + bool is_terminal(size_t s) const { return match[s] >= 0; } + + // index of the longest pattern ending at this state, -1 if none + int32_t match_pattern(size_t s) const { return match[s]; } + + // follow failure links until a transition on `ch` exists. + size_t next(size_t state, uint32_t ch) const; +}; diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp index 371a305fc3..4dd00efddf 100644 --- a/tests/test-chat.cpp +++ b/tests/test-chat.cpp @@ -1144,7 +1144,7 @@ static void test_peg_parser(common_chat_templates * tmpls, // budget sampler inhibits grammar application while inside thinking blocks — // triggers inside ... are suppressed. bool use_reasoning_budget_path = false; - if (parser.params_.grammar_lazy && !parser.params_.thinking_end_tag.empty()) { + if (parser.params_.grammar_lazy && !parser.params_.thinking_end_tags.empty()) { use_reasoning_budget_path = true; for (const auto & trigger : parser.params_.grammar_triggers) { if (trigger.type != COMMON_GRAMMAR_TRIGGER_TYPE_WORD) { @@ -1162,7 +1162,7 @@ static void test_peg_parser(common_chat_templates * tmpls, // Walk through full_input tracking thinking state; only match triggers // when outside thinking blocks. const auto & think_start = parser.params_.thinking_start_tag; - const auto & think_end = parser.params_.thinking_end_tag; + const auto & think_ends = parser.params_.thinking_end_tags; bool in_thinking = false; for (size_t i = 0; i < full_input.size(); ++i) { @@ -1172,12 +1172,14 @@ static void test_peg_parser(common_chat_templates * tmpls, i += think_start.size() - 1; continue; } - if (in_thinking && full_input.compare(i, think_end.size(), think_end) == 0) { - in_thinking = false; - i += think_end.size() - 1; - continue; - } if (in_thinking) { + for (const auto & think_end : think_ends) { + if (full_input.compare(i, think_end.size(), think_end) == 0) { + in_thinking = false; + i += think_end.size() - 1; + break; + } + } continue; } // Outside thinking — check if any trigger word starts here diff --git a/tests/test-reasoning-budget.cpp b/tests/test-reasoning-budget.cpp index f54cff4f8a..3bcc77e173 100644 --- a/tests/test-reasoning-budget.cpp +++ b/tests/test-reasoning-budget.cpp @@ -20,8 +20,8 @@ static void test_reasoning_budget( const char * test_name, const std::vector & sequence, - const std::vector & start_tokens, - const std::vector & end_tokens, + const std::vector & start_seqs, + const std::vector & end_seqs, const std::vector & forced_tokens, int32_t budget, common_reasoning_budget_state initial_state, @@ -31,8 +31,12 @@ static void test_reasoning_budget( // Find the maximum token ID to ensure our vocab covers all tokens llama_token max_token = 0; for (auto t : sequence) max_token = std::max(max_token, t); - for (auto t : start_tokens) max_token = std::max(max_token, t); - for (auto t : end_tokens) max_token = std::max(max_token, t); + for (const auto & seq : start_seqs) { + for (auto t : seq) max_token = std::max(max_token, t); + } + for (const auto & seq : end_seqs) { + for (auto t : seq) max_token = std::max(max_token, t); + } for (auto t : forced_tokens) max_token = std::max(max_token, t); // Create a minimal sampler with mock vocabulary @@ -40,8 +44,8 @@ static void test_reasoning_budget( // The UTF-8 boundary check will treat all tokens as complete (safe fallback) auto * sampler = common_reasoning_budget_init( nullptr, // vocab - not used for basic state machine tests - start_tokens, - end_tokens, + start_seqs, + end_seqs, forced_tokens, budget, initial_state @@ -152,7 +156,7 @@ static void test_reasoning_budget_clone_mid_counting() { const std::vector end = {101}; const std::vector forced = {102, 101}; - auto * sampler = common_reasoning_budget_init(nullptr, start, end, forced, 2, REASONING_BUDGET_IDLE); + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 2, REASONING_BUDGET_IDLE); llama_sampler_accept(sampler, 100); // COUNTING, remaining=2 llama_sampler_accept(sampler, 50); // COUNTING, remaining=1 @@ -171,7 +175,7 @@ static void test_reasoning_budget_clone_mid_forcing() { const std::vector end = {101}; const std::vector forced = {102, 101}; - auto * sampler = common_reasoning_budget_init(nullptr, start, end, forced, 0, REASONING_BUDGET_FORCING); + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 0, REASONING_BUDGET_FORCING); GGML_ASSERT(get_forced_token(sampler, 102) == 102); llama_sampler_accept(sampler, 102); // advance to the second forced token @@ -191,7 +195,7 @@ static void test_reasoning_budget_force_manual() { // if COUNTING, force() succeeds and begins forcing the end sequence from the start { - auto * sampler = common_reasoning_budget_init(nullptr, start, end, forced, 5, REASONING_BUDGET_IDLE); + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 5, REASONING_BUDGET_IDLE); llama_sampler_accept(sampler, 100); // COUNTING, remaining=5 llama_sampler_accept(sampler, 50); // COUNTING, remaining=4 @@ -212,7 +216,7 @@ static void test_reasoning_budget_force_manual() { // if IDLE, force() is a no-op { - auto * sampler = common_reasoning_budget_init(nullptr, start, end, forced, 5, REASONING_BUDGET_IDLE); + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 5, REASONING_BUDGET_IDLE); GGML_ASSERT(!common_reasoning_budget_force(sampler) && "force() must not transition from IDLE"); GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_IDLE); @@ -222,7 +226,7 @@ static void test_reasoning_budget_force_manual() { // if DONE, force() is a no-op { - auto * sampler = common_reasoning_budget_init(nullptr, start, end, forced, 5, REASONING_BUDGET_IDLE); + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 5, REASONING_BUDGET_IDLE); llama_sampler_accept(sampler, 100); // COUNTING llama_sampler_accept(sampler, 101); // natural end -> DONE @@ -236,7 +240,7 @@ static void test_reasoning_budget_force_manual() { // if FORCING, force() is a no-op and must not rewind the force position { - auto * sampler = common_reasoning_budget_init(nullptr, start, end, forced, 0, REASONING_BUDGET_FORCING); + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 0, REASONING_BUDGET_FORCING); GGML_ASSERT(get_forced_token(sampler, 102) == 102); llama_sampler_accept(sampler, 102); // advance to the second forced token (force_pos=1) @@ -254,6 +258,81 @@ static void test_reasoning_budget_force_manual() { fprintf(stderr, " Test 'manual force transition' passed\n"); } +static void test_reasoning_budget_end_match() { + const std::vector start = {{100}}; + const std::vector end = {{101}, {103, 104}}; + + // natural end records the sequence that matched; re-arming clears it + { + auto * sampler = common_reasoning_budget_init(nullptr, start, end, {102, 101}, 5, REASONING_BUDGET_IDLE); + + GGML_ASSERT(common_reasoning_budget_get_end_match(sampler) == nullptr); + + llama_sampler_accept(sampler, 100); // COUNTING + llama_sampler_accept(sampler, 50); + llama_sampler_accept(sampler, 103); + llama_sampler_accept(sampler, 104); // end matched via {103, 104}, DONE + + const llama_tokens * matched = common_reasoning_budget_get_end_match(sampler); + GGML_ASSERT(matched != nullptr); + GGML_ASSERT(*matched == llama_tokens({103, 104})); + + llama_sampler_accept(sampler, 100); // re-arm, COUNTING + GGML_ASSERT(common_reasoning_budget_get_end_match(sampler) == nullptr); + + llama_sampler_free(sampler); + } + + // overlapping end sequences: the longest one ending at the position wins + { + const std::vector end_overlap = {{104}, {103, 104}}; + + auto * sampler = common_reasoning_budget_init(nullptr, start, end_overlap, {102, 104}, 5, REASONING_BUDGET_IDLE); + + llama_sampler_accept(sampler, 100); // COUNTING + llama_sampler_accept(sampler, 103); + llama_sampler_accept(sampler, 104); // both {104} and {103, 104} end here + + const llama_tokens * matched = common_reasoning_budget_get_end_match(sampler); + GGML_ASSERT(matched != nullptr); + GGML_ASSERT(*matched == llama_tokens({103, 104})); + + llama_sampler_free(sampler); + } + + // forcing records the end sequence terminating forced_tokens + { + auto * sampler = common_reasoning_budget_init(nullptr, start, end, {102, 103, 104}, 0, REASONING_BUDGET_FORCING); + + llama_sampler_accept(sampler, 102); + llama_sampler_accept(sampler, 103); + GGML_ASSERT(common_reasoning_budget_get_end_match(sampler) == nullptr); + llama_sampler_accept(sampler, 104); // forced sequence complete, DONE + + const llama_tokens * matched = common_reasoning_budget_get_end_match(sampler); + GGML_ASSERT(matched != nullptr); + GGML_ASSERT(*matched == llama_tokens({103, 104})); + + llama_sampler_free(sampler); + } + + // forced_tokens not ending with a known end sequence records nothing + { + auto * sampler = common_reasoning_budget_init(nullptr, start, end, {102}, 0, REASONING_BUDGET_FORCING); + + llama_sampler_accept(sampler, 102); // forced sequence complete, DONE + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_DONE); + GGML_ASSERT(common_reasoning_budget_get_end_match(sampler) == nullptr); + + llama_sampler_free(sampler); + } + + // a null sampler is safely ignored + GGML_ASSERT(common_reasoning_budget_get_end_match(nullptr) == nullptr); + + fprintf(stderr, " Test 'matched end sequence' passed\n"); +} + // UTF-8 boundary detection unit test // Tests common_utf8_is_complete() from reasoning-budget.h static void test_utf8_boundary_detection() { @@ -290,7 +369,7 @@ int main(void) { const std::vector forced = {102}; // forced token (not used in this test) const std::vector sequence = {100, 50, 51, 101, 52}; // start, two tokens, end, one more - test_reasoning_budget("natural end before budget exhausted", sequence, start, end, forced, + test_reasoning_budget("natural end before budget exhausted", sequence, {start}, {end}, forced, 5, // budget of 5 tokens REASONING_BUDGET_IDLE, SIZE_MAX, SIZE_MAX); // no forcing expected (natural end) @@ -306,7 +385,7 @@ int main(void) { const std::vector forced = {102, 101}; // forced message + end const std::vector sequence = {100, 50, 51, 52, 53}; // start + 4 tokens (budget=2) - test_reasoning_budget("budget exhausted forcing", sequence, start, end, forced, + test_reasoning_budget("budget exhausted forcing", sequence, {start}, {end}, forced, 2, // budget of 2 tokens REASONING_BUDGET_IDLE, 3, // forcing starts at i=3 (accept at i=2 depletes budget, apply at i=3 forces) @@ -321,7 +400,7 @@ int main(void) { const std::vector forced = {102, 101}; const std::vector sequence = {100, 50, 51, 52}; // start token first, then 3 tokens - test_reasoning_budget("activate immediately budget=0", sequence, start, end, forced, + test_reasoning_budget("activate immediately budget=0", sequence, {start}, {end}, forced, 0, // budget of 0 tokens REASONING_BUDGET_COUNTING, // starts counting, promoted to FORCING since budget=0 0, // forcing starts at i=0 (initialized in FORCING, apply forces immediately) @@ -335,7 +414,7 @@ int main(void) { const std::vector forced = {102}; const std::vector sequence = {50, 51, 52, 53}; - test_reasoning_budget("no start/end configured", sequence, start, end, forced, + test_reasoning_budget("no start/end configured", sequence, {start}, {end}, forced, 2, // budget REASONING_BUDGET_IDLE, SIZE_MAX, SIZE_MAX); // no forcing (no start/end configured) @@ -350,7 +429,7 @@ int main(void) { const std::vector forced = {102, 101}; const std::vector sequence = {50, 51, 52, 53}; - test_reasoning_budget("activate immediately with budget", sequence, start, end, forced, + test_reasoning_budget("activate immediately with budget", sequence, {start}, {end}, forced, 2, // budget of 2 tokens REASONING_BUDGET_COUNTING, 2, // forcing starts at i=2 (after 2 accepts deplete budget, apply at i=2 forces) @@ -373,18 +452,50 @@ int main(void) { const std::vector forced = {102, 101}; const std::vector sequence = {100, 50, 101, 100, 60, 61, 62, 63}; - test_reasoning_budget("multi-block re-arms budget after DONE", sequence, start, end, forced, + test_reasoning_budget("multi-block re-arms budget after DONE", sequence, {start}, {end}, forced, 2, // budget of 2 tokens (per block) REASONING_BUDGET_IDLE, 6, // forcing starts at i=6 (after second block exhausts at i=5) 7); // forcing continues through i=7 } + // Test 7: Multiple start sequences - the second sequence activates counting + // Flow: i=0 accept(110), i=1 accept(111)->COUNTING rem=2; i=2 accept(50)->rem=1; + // i=3 accept(51)->rem=0->FORCING; i=4..5 apply() forces the end sequence + { + const std::vector start = {{100}, {110, 111}}; + const std::vector end = {{101}}; + const std::vector forced = {102, 101}; + const std::vector sequence = {110, 111, 50, 51, 52, 53}; + + test_reasoning_budget("multiple start sequences", sequence, start, end, forced, + 2, // budget of 2 tokens + REASONING_BUDGET_IDLE, + 4, // forcing starts at i=4 (accept at i=3 depletes budget) + 5); // forcing continues through i=5 + } + + // Test 8: Multiple end sequences - natural end via the second sequence + // Flow: i=0 accept(100)->COUNTING rem=5; i=1 accept(50)->rem=4; + // i=2 accept(103)->partial end, rem=3; i=3 accept(104)->end matched, DONE + { + const std::vector start = {{100}}; + const std::vector end = {{101}, {103, 104}}; + const std::vector forced = {102, 101}; + const std::vector sequence = {100, 50, 103, 104, 52}; + + test_reasoning_budget("multiple end sequences", sequence, start, end, forced, + 5, // budget of 5 tokens + REASONING_BUDGET_IDLE, + SIZE_MAX, SIZE_MAX); // no forcing expected (natural end) + } + test_reasoning_budget_clone_mid_counting(); test_reasoning_budget_clone_mid_forcing(); test_reasoning_budget_force_manual(); + test_reasoning_budget_end_match(); - printf("OK (9 tests passed)\n"); + printf("OK (12 tests passed)\n"); printf("Testing UTF-8 boundary detection... "); test_utf8_boundary_detection(); diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp index da58049308..c9109fc962 100644 --- a/tools/server/server-common.cpp +++ b/tools/server/server-common.cpp @@ -1131,10 +1131,10 @@ json oaicompat_chat_params_parse( reasoning_budget = opt.reasoning_budget; } - if (!chat_params.thinking_end_tag.empty()) { + if (!chat_params.thinking_end_tags.empty()) { llama_params["reasoning_budget_tokens"] = reasoning_budget; llama_params["reasoning_budget_start_tag"] = chat_params.thinking_start_tag; - llama_params["reasoning_budget_end_tag"] = chat_params.thinking_end_tag; + llama_params["reasoning_budget_end_tags"] = chat_params.thinking_end_tags; llama_params["reasoning_budget_message"] = json_value(body, "reasoning_budget_message", opt.reasoning_budget_message); llama_params["reasoning_control"] = json_value(body, "reasoning_control", false); } diff --git a/tools/server/server-schema.cpp b/tools/server/server-schema.cpp index 89026eb4e3..e880f4ca72 100644 --- a/tools/server/server-schema.cpp +++ b/tools/server/server-schema.cpp @@ -390,21 +390,40 @@ std::vector> make_llama_cmpl_schema(const common_params & ctx.params.sampling.reasoning_budget_start = common_tokenize(ctx.vocab, data.at("reasoning_budget_start_tag").get(), false, true); })); - add((new field_str("reasoning_budget_end_tag")) - ->set_desc("Token string marking the end of the reasoning budget section") + add((new field_json("reasoning_budget_end_tags")) + ->add_alias("reasoning_budget_end_tag") + ->set_desc("Token strings marking the end of the reasoning budget section; the first is forced when the budget expires") ->set_handler([&](field_eval_context & ctx, const json & data) { GGML_ASSERT(ctx.vocab != nullptr); - std::string end_tag = data.at("reasoning_budget_end_tag").get(); - ctx.params.sampling.reasoning_budget_end = common_tokenize(ctx.vocab, end_tag, false, true); + ctx.params.sampling.reasoning_budget_end.clear(); + if (data.contains("reasoning_budget_end_tags")) { + for (const auto & t : data.at("reasoning_budget_end_tags")) { + std::string tag = t.get(); + if (!tag.empty()) { + ctx.params.sampling.reasoning_budget_end.push_back(common_tokenize(ctx.vocab, tag, false, true)); + } + } + } else if (data.contains("reasoning_budget_end_tag")) { + std::string tag = data.at("reasoning_budget_end_tag").get(); + if (!tag.empty()) { + ctx.params.sampling.reasoning_budget_end.push_back(common_tokenize(ctx.vocab, tag, false, true)); + } + } })); add((new field_str("reasoning_budget_message")) ->set_desc("Message to prepend to the reasoning budget end tag when forcing it") ->set_handler([&](field_eval_context & ctx, const json & data) { GGML_ASSERT(ctx.vocab != nullptr); - std::string end_tag = json_value(data, "reasoning_budget_end_tag", std::string()); - std::string message = data.at("reasoning_budget_message").get(); - ctx.params.sampling.reasoning_budget_forced = common_tokenize(ctx.vocab, message + end_tag, false, true); + if (!ctx.params.sampling.reasoning_budget_end.empty()) { + llama_tokens end_tag = ctx.params.sampling.reasoning_budget_end.front(); + std::string message = json_value(data, "reasoning_budget_message", std::string()); + if (!message.empty()) { + llama_tokens message_tokens = common_tokenize(ctx.vocab, message, false, true); + end_tag.insert(end_tag.begin(), message_tokens.begin(), message_tokens.end()); + } + ctx.params.sampling.reasoning_budget_forced = std::move(end_tag); + } })); add((new field_json("logit_bias")) @@ -546,7 +565,7 @@ task_params eval_llama_cmpl_schema( // debugging { auto budget = params.sampling.reasoning_budget_tokens; - SRV_DBG("reasoning budget: tokens=%d, generation_prompt='%s', start=%zu toks, end=%zu toks, forced=%zu toks\n", + SRV_DBG("reasoning budget: tokens=%d, generation_prompt='%s', start=%zu toks, end=%zu seqs, forced=%zu toks\n", budget, params.sampling.generation_prompt.c_str(), params.sampling.reasoning_budget_start.size(), params.sampling.reasoning_budget_end.size(), From fb92d8f1873c96ec63f9c59721d58a55bf46d441 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Yongmin=20Yoo=20=EC=9C=A0=EC=9A=A9=EB=AF=BC?= Date: Sat, 25 Jul 2026 21:32:37 +0900 Subject: [PATCH 3/9] Update ggml/src/gguf.cpp : Defined virtual keyword for destructor of gguf_writer_base (#25867) Without a virtual destructor, deleting a derived object through a base-class pointer only invokes the base destructor, skipping the derived one. --- ggml/src/gguf.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ggml/src/gguf.cpp b/ggml/src/gguf.cpp index 7920b8100b..9f9e4fe5d1 100644 --- a/ggml/src/gguf.cpp +++ b/ggml/src/gguf.cpp @@ -1424,7 +1424,7 @@ void gguf_set_tensor_data(struct gguf_context * ctx, const char * name, const vo struct gguf_writer_base { size_t written_bytes {0u}; - ~gguf_writer_base(void) = default; + virtual ~gguf_writer_base(void) = default; // we bet on devirtualization virtual void write(int8_t val) = 0; From 720d7fa4097f76e5d0eade5a92c1df87c1faf9d9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adrien=20Gallou=C3=ABt?= Date: Sat, 25 Jul 2026 18:16:29 +0200 Subject: [PATCH 4/9] vendor : update cpp-httplib to 0.51.0 (#26067) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Adrien Gallouët --- scripts/sync_vendor.py | 2 +- vendor/cpp-httplib/httplib.cpp | 321 ++++++++++++++++++++++++++------- vendor/cpp-httplib/httplib.h | 47 ++++- 3 files changed, 296 insertions(+), 74 deletions(-) diff --git a/scripts/sync_vendor.py b/scripts/sync_vendor.py index a9343e37c1..679e557a32 100755 --- a/scripts/sync_vendor.py +++ b/scripts/sync_vendor.py @@ -5,7 +5,7 @@ import os import sys import subprocess -HTTPLIB_VERSION = "refs/tags/v0.50.1" +HTTPLIB_VERSION = "refs/tags/v0.51.0" vendor = { "https://github.com/nlohmann/json/releases/latest/download/json.hpp": "vendor/nlohmann/json.hpp", diff --git a/vendor/cpp-httplib/httplib.cpp b/vendor/cpp-httplib/httplib.cpp index e58fb3eea6..1551796387 100644 --- a/vendor/cpp-httplib/httplib.cpp +++ b/vendor/cpp-httplib/httplib.cpp @@ -3161,9 +3161,7 @@ get_multimap_value(const Map &m, const std::string &key, size_t id) { void set_header(Headers &headers, const std::string &key, const std::string &val) { - if (fields::is_field_name(key) && fields::is_field_value(val)) { - headers.emplace(key, val); - } + if (fields::is_field_valid(key, val)) { headers.emplace(key, val); } } bool read_headers(Stream &strm, Headers &headers) { @@ -3370,8 +3368,46 @@ ReadContentResult read_content_chunked(Stream &strm, T &x, } bool is_chunked_transfer_encoding(const Headers &headers) { - return case_ignore::equal( - get_header_value(headers, "Transfer-Encoding", "", 0), "chunked"); + // RFC 9112 6.1: a message is framed with the chunked coding when "chunked" + // is the final transfer coding. A single field value may list several + // codings ("gzip, chunked"), and the list may be split across multiple + // Transfer-Encoding header lines (RFC 9110 5.3). Match the last coding token + // case-insensitively rather than comparing the whole value against "chunked". + // + // Security: reading a chunked message as unframed leaves its body in the + // socket, where a keep-alive connection parses it as a smuggled request. + // Headers is an unordered_multimap whose iteration order for duplicate keys + // is not portable, so when there is more than one Transfer-Encoding line we + // cannot tell which coding is truly final. In that ambiguous case we fail + // safe by treating the message as chunked (a mis-parse just closes the + // connection, whereas the opposite error enables smuggling). + auto rng = headers.equal_range("Transfer-Encoding"); + + size_t line_count = 0; + bool chunked_present = false; + bool last_line_ends_with_chunked = false; + + for (auto it = rng.first; it != rng.second; ++it) { + line_count++; + const auto &value = it->second; + + std::string last_coding; + bool line_has_chunked = false; + split(value.data(), value.data() + value.size(), ',', + [&](const char *b, const char *e) { + last_coding.assign(b, e); + if (case_ignore::equal(last_coding, "chunked")) { + line_has_chunked = true; + } + }); + + if (line_has_chunked) { chunked_present = true; } + last_line_ends_with_chunked = case_ignore::equal(last_coding, "chunked"); + } + + if (line_count == 0) { return false; } + if (line_count == 1) { return last_line_ends_with_chunked; } + return chunked_present; } template @@ -3488,6 +3524,13 @@ bool read_content(Stream &strm, T &x, size_t payload_max_length, int &status, ssize_t write_request_line(Stream &strm, const std::string &method, const std::string &path) { + // A request target must not carry CR/LF (or other control octets); otherwise + // a value smuggled into it splits the request line and injects headers or a + // whole request. The same field-value check already guards header values in + // check_and_write_headers and the request target in + // perform_websocket_handshake; apply it here too. + if (!fields::is_field_value(path)) { return -1; } + std::string s = method; s += ' '; s += path; @@ -3507,6 +3550,13 @@ ssize_t write_response_line(Stream &strm, int status) { ssize_t write_headers(Stream &strm, const Headers &headers) { ssize_t write_len = 0; for (const auto &x : headers) { + // Skip fields with invalid names or values to prevent response splitting + // via CR/LF injection, matching set_header(). The client validates request + // headers up front in check_and_write_headers, but the server passes + // res.headers straight to this writer, and res.headers is a public field + // an application can populate directly with request-derived values. + if (!fields::is_field_valid(x.first, x.second)) { continue; } + std::string s; s = x.first; s += ": "; @@ -3707,10 +3757,7 @@ write_content_chunked(Stream &strm, const ContentProvider &content_provider, for (const auto &kv : *trailer) { // Skip fields with invalid names or values to prevent response // splitting via CR/LF injection, matching set_header(). - if (!fields::is_field_name(kv.first) || - !fields::is_field_value(kv.second)) { - continue; - } + if (!fields::is_field_valid(kv.first, kv.second)) { continue; } std::string field_line = kv.first + ": " + kv.second + "\r\n"; if (!write_data(strm, field_line.data(), field_line.size())) { ok = false; @@ -4079,6 +4126,13 @@ public: break; } + // Check header count limit + if (header_count_ >= CPPHTTPLIB_HEADER_MAX_COUNT) { + is_valid_ = false; + return false; + } + header_count_++; + const auto header = buf_head(pos); if (!parse_header(header.data(), header.data() + header.size(), @@ -4194,6 +4248,7 @@ private: file_.filename.clear(); file_.content_type.clear(); file_.headers.clear(); + header_count_ = 0; } bool start_with_case_ignore(const std::string &a, const char *b, @@ -4247,6 +4302,7 @@ private: size_t state_ = 0; bool is_valid_ = false; FormData file_; + size_t header_count_ = 0; // Buffer bool start_with(const std::string &a, size_t spos, size_t epos, @@ -4907,6 +4963,10 @@ bool is_field_content(const std::string &s) { bool is_field_value(const std::string &s) { return is_field_content(s); } +bool is_field_valid(const std::string &name, const std::string &value) { + return is_field_name(name) && is_field_value(value); +} + } // namespace fields bool perform_websocket_handshake(Stream &strm, const std::string &host, @@ -4921,9 +4981,7 @@ bool perform_websocket_handshake(Stream &strm, const std::string &host, // Validate user-provided headers for (const auto &h : headers) { - if (!fields::is_field_name(h.first) || !fields::is_field_value(h.second)) { - return false; - } + if (!fields::is_field_valid(h.first, h.second)) { return false; } } // Generate random Sec-WebSocket-Key @@ -5042,9 +5100,31 @@ std::string hash_to_hex(const unsigned char (&hash)[N]) { } } // namespace +#ifdef CPPHTTPLIB_MBEDTLS_V4 +// Mbed TLS 4.x provides hashing (and TLS RNG) via PSA Crypto, which must be +// initialized once. PSA state is process-global; do not free it. +bool ensure_mbedtls_psa_crypto() { + static std::once_flag once; + static bool ok = false; + std::call_once(once, []() { ok = (psa_crypto_init() == PSA_SUCCESS); }); + return ok; +} + +bool psa_hash(psa_algorithm_t alg, const std::string &s, + unsigned char *out, size_t out_size) { + if (!ensure_mbedtls_psa_crypto()) { return false; } + size_t olen = 0; + return psa_hash_compute(alg, reinterpret_cast(s.data()), + s.size(), out, out_size, &olen) == PSA_SUCCESS && + olen == out_size; +} +#endif + std::string MD5(const std::string &s) { unsigned char hash[16]; -#ifdef CPPHTTPLIB_MBEDTLS_V3 +#ifdef CPPHTTPLIB_MBEDTLS_V4 + if (!psa_hash(PSA_ALG_MD5, s, hash, sizeof(hash))) { return {}; } +#elif defined(CPPHTTPLIB_MBEDTLS_V3) mbedtls_md5(reinterpret_cast(s.c_str()), s.size(), hash); #else @@ -5056,7 +5136,9 @@ std::string MD5(const std::string &s) { std::string SHA_256(const std::string &s) { unsigned char hash[32]; -#ifdef CPPHTTPLIB_MBEDTLS_V3 +#ifdef CPPHTTPLIB_MBEDTLS_V4 + if (!psa_hash(PSA_ALG_SHA_256, s, hash, sizeof(hash))) { return {}; } +#elif defined(CPPHTTPLIB_MBEDTLS_V3) mbedtls_sha256(reinterpret_cast(s.c_str()), s.size(), hash, 0); #else @@ -5068,7 +5150,9 @@ std::string SHA_256(const std::string &s) { std::string SHA_512(const std::string &s) { unsigned char hash[64]; -#ifdef CPPHTTPLIB_MBEDTLS_V3 +#ifdef CPPHTTPLIB_MBEDTLS_V4 + if (!psa_hash(PSA_ALG_SHA_512, s, hash, sizeof(hash))) { return {}; } +#elif defined(CPPHTTPLIB_MBEDTLS_V3) mbedtls_sha512(reinterpret_cast(s.c_str()), s.size(), hash, 0); #else @@ -6949,8 +7033,7 @@ template bool check_and_write_headers(Stream &strm, Headers &headers, T header_writer, Error &error) { for (const auto &h : headers) { - if (!detail::fields::is_field_name(h.first) || - !detail::fields::is_field_value(h.second)) { + if (!detail::fields::is_field_valid(h.first, h.second)) { error = Error::InvalidHeaders; return false; } @@ -8348,25 +8431,25 @@ get_client_ip(const std::string &x_forwarded_for, // caller can fall back to the connection-level remote address. if (ip_list.empty()) { return std::string(); } - for (size_t i = 0; i < ip_list.size(); ++i) { - auto ip = ip_list[i]; + // Each hop appends the address it received the request from, so the rightmost + // entries are the ones written by our own infrastructure while the leftmost + // are whatever the original client chose to send. Walk from the right and + // skip trusted proxies; the first address that is not a trusted proxy is the + // furthest point still attributable to a real hop, i.e. the client. Scanning + // from the left instead lets a client forge an arbitrary address by following + // it with a trusted proxy's address, which the left-to-right scan then + // returned as the client. + for (size_t i = ip_list.size(); i-- > 0;) { + const auto &ip = ip_list[i]; auto is_trusted_proxy = std::any_of(trusted_proxies.begin(), trusted_proxies.end(), [&](const std::string &proxy) { return ip == proxy; }); - if (is_trusted_proxy) { - if (i == 0) { - // If the trusted proxy is the first IP, there's no preceding client IP - return ip; - } else { - // Return the IP immediately before the trusted proxy - return ip_list[i - 1]; - } - } + if (!is_trusted_proxy) { return ip; } } - // If no trusted proxy is found, return the first IP in the list + // Every hop was a trusted proxy; fall back to the first entry. return ip_list.front(); } @@ -8436,7 +8519,14 @@ Server::process_request(Stream &strm, const std::string &remote_addr, connection_closed = true; } - if (!trusted_proxies_.empty() && req.has_header("X-Forwarded-For")) { + // Only honor X-Forwarded-For if the peer on the actual TCP connection is + // itself a trusted proxy. Otherwise any direct client could spoof + // remote_addr simply by sending an arbitrary X-Forwarded-For header. + auto is_trusted_peer = std::any_of( + trusted_proxies_.begin(), trusted_proxies_.end(), + [&](const std::string &proxy) { return proxy == remote_addr; }); + + if (is_trusted_peer && req.has_header("X-Forwarded-For")) { auto x_forwarded_for = req.get_header_value("X-Forwarded-For"); auto derived = get_client_ip(x_forwarded_for, trusted_proxies_); req.remote_addr = derived.empty() ? remote_addr : derived; @@ -8643,13 +8733,19 @@ Server::process_request(Stream &strm, const std::string &remote_addr, // Drain any unconsumed framed body to prevent request smuggling on // keep-alive. Without framing there is no body to drain — reading would - // consume the next request (issue #2450). + // consume the next request (issue #2450). If the response has committed the + // connection to close, there is no next request to protect. if (!req.body_consumed_ && detail::has_framed_body(req)) { - int dummy_status; - if (!detail::read_content( - strm, req, payload_max_length_, dummy_status, nullptr, - [](const char *, size_t, size_t, size_t) { return true; }, false)) { + if (res.get_header_value("Connection") == "close") { connection_closed = true; + } else { + int dummy_status; + if (!detail::read_content( + strm, req, payload_max_length_, dummy_status, nullptr, + [](const char *, size_t, size_t, size_t) { return true; }, + false)) { + connection_closed = true; + } } } @@ -9210,9 +9306,8 @@ ClientImpl::open_stream(const std::string &method, const std::string &path, handle.body_reader_.content_length = content_length; } - auto transfer_encoding = - handle.response->get_header_value("Transfer-Encoding"); - handle.body_reader_.chunked = (transfer_encoding == "chunked"); + handle.body_reader_.chunked = + detail::is_chunked_transfer_encoding(handle.response->headers); auto content_encoding = handle.response->get_header_value("Content-Encoding"); if (!content_encoding.empty()) { @@ -9541,8 +9636,8 @@ bool ClientImpl::create_redirect_client( // Clean up request headers that are host/client specific // Remove headers that should not be carried over to new host - auto headers_to_remove = - std::vector{"Host", "Proxy-Authorization", "Authorization"}; + auto headers_to_remove = std::vector{ + "Host", "Proxy-Authorization", "Authorization", "Cookie", "Cookie2"}; for (const auto &header_name : headers_to_remove) { auto it = req.headers.find(header_name); @@ -9796,7 +9891,14 @@ bool ClientImpl::write_request(Stream &strm, Request &req, } // Write request line and headers - detail::write_request_line(bstrm, req.method, path_with_query); + if (detail::write_request_line(bstrm, req.method, path_with_query) < 0) { + // A rejected target (e.g. CR/LF smuggled in via a decoded redirect + // Location under set_path_encode(false)) must fail the request cleanly + // instead of emitting a request-line-less, header-injecting request. + error = Error::Write; + output_error_log(error, &req); + return false; + } if (!detail::check_and_write_headers(bstrm, req.headers, header_writer_, error)) { output_error_log(error, &req); @@ -13893,6 +13995,13 @@ struct MbedTlsSession { std::string hostname; // For client: set via set_sni std::string sni_hostname; // For server: received from client via SNI callback + // Mbed TLS has no SSL_peek() equivalent, so is_peer_closed() must probe with + // a real 1-byte mbedtls_ssl_read(). If that probe lands on application data + // (e.g. a response that arrived while this side was still in its post-write + // check), the byte is pushed back here and served by the next read(). + unsigned char peeked_byte = 0; + bool has_peeked_byte = false; + MbedTlsSession() { mbedtls_ssl_init(&ssl); } ~MbedTlsSession() { mbedtls_ssl_free(&ssl); } @@ -13927,6 +14036,20 @@ ErrorCode map_mbedtls_error(int ret, int &out_errno) { return ErrorCode::Fatal; } +// A TLS 1.3 NewSessionTicket (signaled by default on Mbed TLS 4.x) is a +// non-fatal notification delivered between records, not an error and not +// application data, so I/O calls that see it should just be retried. Kept in +// one helper so the retry loops keep an intact "do { } while (...)" instead of +// splitting the closing brace across an #if. +bool mbedtls_is_session_ticket(int ret) { +#if defined(MBEDTLS_ERR_SSL_RECEIVED_NEW_SESSION_TICKET) + return ret == MBEDTLS_ERR_SSL_RECEIVED_NEW_SESSION_TICKET; +#else + (void)ret; + return false; +#endif +} + // BIO-like send callback for Mbed TLS int mbedtls_net_send_cb(void *ctx, const unsigned char *buf, size_t len) { @@ -13978,8 +14101,10 @@ int mbedtls_net_recv_cb(void *ctx, unsigned char *buf, size_t len) { // MbedTlsContext constructor/destructor implementations MbedTlsContext::MbedTlsContext() { mbedtls_ssl_config_init(&conf); +#ifndef CPPHTTPLIB_MBEDTLS_V4 mbedtls_entropy_init(&entropy); mbedtls_ctr_drbg_init(&ctr_drbg); +#endif mbedtls_x509_crt_init(&ca_chain); mbedtls_x509_crt_init(&own_cert); mbedtls_pk_init(&own_key); @@ -13989,8 +14114,10 @@ MbedTlsContext::~MbedTlsContext() { mbedtls_pk_free(&own_key); mbedtls_x509_crt_free(&own_cert); mbedtls_x509_crt_free(&ca_chain); +#ifndef CPPHTTPLIB_MBEDTLS_V4 mbedtls_ctr_drbg_free(&ctr_drbg); mbedtls_entropy_free(&entropy); +#endif mbedtls_ssl_config_free(&conf); } @@ -14064,6 +14191,14 @@ ctx_t create_client_context() { ctx->is_server = false; +#ifdef CPPHTTPLIB_MBEDTLS_V4 + // Mbed TLS 4.x draws randomness from PSA Crypto; just ensure it is ready. + if (!detail::ensure_mbedtls_psa_crypto()) { + delete ctx; + return nullptr; + } + int ret; +#else // Seed the random number generator const char *pers = "httplib_client"; int ret = mbedtls_ctr_drbg_seed( @@ -14074,6 +14209,7 @@ ctx_t create_client_context() { delete ctx; return nullptr; } +#endif // Set up SSL config for client ret = mbedtls_ssl_config_defaults(&ctx->conf, MBEDTLS_SSL_IS_CLIENT, @@ -14085,8 +14221,10 @@ ctx_t create_client_context() { return nullptr; } - // Set random number generator +#ifndef CPPHTTPLIB_MBEDTLS_V4 + // Set random number generator (Mbed TLS 4.x uses the PSA RNG implicitly) mbedtls_ssl_conf_rng(&ctx->conf, mbedtls_ctr_drbg_random, &ctx->ctr_drbg); +#endif // Default: verify peer certificate mbedtls_ssl_conf_authmode(&ctx->conf, MBEDTLS_SSL_VERIFY_REQUIRED); @@ -14108,6 +14246,14 @@ ctx_t create_server_context() { ctx->is_server = true; +#ifdef CPPHTTPLIB_MBEDTLS_V4 + // Mbed TLS 4.x draws randomness from PSA Crypto; just ensure it is ready. + if (!detail::ensure_mbedtls_psa_crypto()) { + delete ctx; + return nullptr; + } + int ret; +#else // Seed the random number generator const char *pers = "httplib_server"; int ret = mbedtls_ctr_drbg_seed( @@ -14118,6 +14264,7 @@ ctx_t create_server_context() { delete ctx; return nullptr; } +#endif // Set up SSL config for server ret = mbedtls_ssl_config_defaults(&ctx->conf, MBEDTLS_SSL_IS_SERVER, @@ -14129,8 +14276,10 @@ ctx_t create_server_context() { return nullptr; } - // Set random number generator +#ifndef CPPHTTPLIB_MBEDTLS_V4 + // Set random number generator (Mbed TLS 4.x uses the PSA RNG implicitly) mbedtls_ssl_conf_rng(&ctx->conf, mbedtls_ctr_drbg_random, &ctx->ctr_drbg); +#endif // Default: don't verify client mbedtls_ssl_conf_authmode(&ctx->conf, MBEDTLS_SSL_VERIFY_NONE); @@ -14290,7 +14439,7 @@ bool set_client_cert_pem(ctx_t ctx, const char *cert, const char *key, password ? reinterpret_cast(password) : nullptr; size_t pwd_len = password ? strlen(password) : 0; -#ifdef CPPHTTPLIB_MBEDTLS_V3 +#if defined(CPPHTTPLIB_MBEDTLS_V3) && !defined(CPPHTTPLIB_MBEDTLS_V4) ret = mbedtls_pk_parse_key( &mctx->own_key, reinterpret_cast(key_str.c_str()), key_str.size() + 1, pwd, pwd_len, mbedtls_ctr_drbg_random, @@ -14305,7 +14454,10 @@ bool set_client_cert_pem(ctx_t ctx, const char *cert, const char *key, return false; } - // Verify that the certificate and private key match + // Verify that the certificate and private key match. + // Mbed TLS 4.x: mbedtls_pk_check_pair() reports a spurious mismatch for + // PSA-backed keys, so skip it and let the handshake surface a real mismatch. +#ifndef CPPHTTPLIB_MBEDTLS_V4 #ifdef CPPHTTPLIB_MBEDTLS_V3 ret = mbedtls_pk_check_pair(&mctx->own_cert.pk, &mctx->own_key, mbedtls_ctr_drbg_random, &mctx->ctr_drbg); @@ -14316,6 +14468,7 @@ bool set_client_cert_pem(ctx_t ctx, const char *cert, const char *key, impl::mbedtls_last_error() = ret; return false; } +#endif ret = mbedtls_ssl_conf_own_cert(&mctx->conf, &mctx->own_cert, &mctx->own_key); if (ret != 0) { @@ -14339,7 +14492,7 @@ bool set_client_cert_file(ctx_t ctx, const char *cert_path, } // Parse private key file -#ifdef CPPHTTPLIB_MBEDTLS_V3 +#if defined(CPPHTTPLIB_MBEDTLS_V3) && !defined(CPPHTTPLIB_MBEDTLS_V4) ret = mbedtls_pk_parse_keyfile(&mctx->own_key, key_path, password, mbedtls_ctr_drbg_random, &mctx->ctr_drbg); #else @@ -14350,7 +14503,9 @@ bool set_client_cert_file(ctx_t ctx, const char *cert_path, return false; } - // Verify that the certificate and private key match + // Verify that the certificate and private key match. + // Mbed TLS 4.x: see set_client_cert() — skip the spurious check_pair. +#ifndef CPPHTTPLIB_MBEDTLS_V4 #ifdef CPPHTTPLIB_MBEDTLS_V3 ret = mbedtls_pk_check_pair(&mctx->own_cert.pk, &mctx->own_key, mbedtls_ctr_drbg_random, &mctx->ctr_drbg); @@ -14361,6 +14516,7 @@ bool set_client_cert_file(ctx_t ctx, const char *cert_path, impl::mbedtls_last_error() = ret; return false; } +#endif ret = mbedtls_ssl_conf_own_cert(&mctx->conf, &mctx->own_cert, &mctx->own_key); if (ret != 0) { @@ -14455,7 +14611,10 @@ TlsError connect(session_t session) { } auto msession = static_cast(session); - int ret = mbedtls_ssl_handshake(&msession->ssl); + int ret; + do { + ret = mbedtls_ssl_handshake(&msession->ssl); + } while (impl::mbedtls_is_session_ticket(ret)); if (ret == 0) { err.code = ErrorCode::Success; @@ -14499,6 +14658,8 @@ bool connect_nonblocking(session_t session, socket_t sock, int ret; while ((ret = mbedtls_ssl_handshake(&msession->ssl)) != 0) { + // Non-fatal TLS 1.3 ticket; retry immediately. + if (impl::mbedtls_is_session_ticket(ret)) { continue; } if (ret == MBEDTLS_ERR_SSL_WANT_READ) { if (detail::select_read(sock, timeout_sec, timeout_usec) > 0) { continue; @@ -14546,8 +14707,28 @@ ssize_t read(session_t session, void *buf, size_t len, TlsError &err) { } auto msession = static_cast(session); - int ret = - mbedtls_ssl_read(&msession->ssl, static_cast(buf), len); + + // Serve a byte consumed by the is_peer_closed() probe before reading more. + if (msession->has_peeked_byte) { + if (len == 0) { return 0; } + auto p = static_cast(buf); + p[0] = msession->peeked_byte; + msession->has_peeked_byte = false; + size_t n = 1; + // Top up with any already-decrypted bytes without risking a block. + if (len > 1 && mbedtls_ssl_get_bytes_avail(&msession->ssl) > 0) { + int extra = mbedtls_ssl_read(&msession->ssl, p + 1, len - 1); + if (extra > 0) { n += static_cast(extra); } + } + err.code = ErrorCode::Success; + return static_cast(n); + } + + int ret; + do { + ret = mbedtls_ssl_read(&msession->ssl, static_cast(buf), + len); + } while (impl::mbedtls_is_session_ticket(ret)); if (ret > 0) { err.code = ErrorCode::Success; @@ -14576,8 +14757,11 @@ ssize_t write(session_t session, const void *buf, size_t len, } auto msession = static_cast(session); - int ret = mbedtls_ssl_write(&msession->ssl, - static_cast(buf), len); + int ret; + do { + ret = mbedtls_ssl_write(&msession->ssl, + static_cast(buf), len); + } while (impl::mbedtls_is_session_ticket(ret)); if (ret > 0) { err.code = ErrorCode::Success; @@ -14599,7 +14783,8 @@ int pending(const_session_t session) { if (!session) { return 0; } auto msession = static_cast(const_cast(session)); - return static_cast(mbedtls_ssl_get_bytes_avail(&msession->ssl)); + return static_cast(mbedtls_ssl_get_bytes_avail(&msession->ssl)) + + (msession->has_peeked_byte ? 1 : 0); } void shutdown(session_t session, bool graceful) { @@ -14625,24 +14810,34 @@ bool is_peer_closed(session_t session, socket_t sock) { if (!session || sock == INVALID_SOCKET) { return true; } auto msession = static_cast(session); - // Check if there's already decrypted data available in the TLS buffer - // If so, the connection is definitely alive - if (mbedtls_ssl_get_bytes_avail(&msession->ssl) > 0) { return false; } + // Check if there's already decrypted or pushed-back data available. + // If so, the connection is definitely alive. + if (msession->has_peeked_byte || + mbedtls_ssl_get_bytes_avail(&msession->ssl) > 0) { + return false; + } // Set socket to non-blocking to avoid blocking on read detail::set_nonblocking(sock, true); auto cleanup = detail::scope_exit([&]() { detail::set_nonblocking(sock, false); }); - // Try a 1-byte read to check connection status - // Note: This will consume the byte if data is available, but for the - // purpose of checking if peer is closed, this should be acceptable - // since we're only called when we expect the connection might be closing + // Probe with a 1-byte read (Mbed TLS has no peek API). If the probe lands + // on application data — e.g. a response that already arrived — push the + // byte back so the next read() delivers it instead of losing it. unsigned char buf; - int ret = mbedtls_ssl_read(&msession->ssl, &buf, 1); + int ret; + do { + ret = mbedtls_ssl_read(&msession->ssl, &buf, 1); + } while (impl::mbedtls_is_session_ticket(ret)); // If we got data or WANT_READ (would block), connection is alive - if (ret > 0 || ret == MBEDTLS_ERR_SSL_WANT_READ) { return false; } + if (ret > 0) { + msession->peeked_byte = buf; + msession->has_peeked_byte = true; + return false; + } + if (ret == MBEDTLS_ERR_SSL_WANT_READ) { return false; } // If we get a peer close notify or a connection reset, the peer is closed return ret == MBEDTLS_ERR_SSL_PEER_CLOSE_NOTIFY || @@ -15049,7 +15244,7 @@ bool update_server_cert(ctx_t ctx, const char *cert_pem, } // Parse private key PEM -#ifdef CPPHTTPLIB_MBEDTLS_V3 +#if defined(CPPHTTPLIB_MBEDTLS_V3) && !defined(CPPHTTPLIB_MBEDTLS_V4) ret = mbedtls_pk_parse_key( &mbed_ctx->own_key, reinterpret_cast(key_pem), strlen(key_pem) + 1, diff --git a/vendor/cpp-httplib/httplib.h b/vendor/cpp-httplib/httplib.h index de026a2532..6df357d706 100644 --- a/vendor/cpp-httplib/httplib.h +++ b/vendor/cpp-httplib/httplib.h @@ -8,8 +8,8 @@ #ifndef CPPHTTPLIB_HTTPLIB_H #define CPPHTTPLIB_HTTPLIB_H -#define CPPHTTPLIB_VERSION "0.50.1" -#define CPPHTTPLIB_VERSION_NUM "0x003201" +#define CPPHTTPLIB_VERSION "0.51.0" +#define CPPHTTPLIB_VERSION_NUM "0x003300" #ifdef _WIN32 #if defined(_WIN32_WINNT) && _WIN32_WINNT < 0x0A00 @@ -420,18 +420,26 @@ using socket_t = int; #endif // CPPHTTPLIB_OPENSSL_SUPPORT #ifdef CPPHTTPLIB_MBEDTLS_SUPPORT -#include -#include +// version.h defines MBEDTLS_VERSION_MAJOR (on 2.x/3.x/4.x alike); it is pulled +// in with this first include group so the version gating below can use it. #include -#include #include #include #include +#include +#include +#include +#if MBEDTLS_VERSION_MAJOR >= 4 +// Mbed TLS 4.x moved hashing/RNG to PSA Crypto and removed these headers. +#include +#else +#include +#include +#include #include #include #include -#include -#include +#endif #ifdef _WIN32 #include #ifdef _MSC_VER @@ -444,7 +452,11 @@ using socket_t = int; #endif #endif -// Mbed TLS 3.x API compatibility +// Mbed TLS version API compatibility. Note: V4 implies V3 (both defined on +// 4.x), so version-specific 3.x-only code must check V3 && !V4. +#if MBEDTLS_VERSION_MAJOR >= 4 +#define CPPHTTPLIB_MBEDTLS_V4 +#endif #if MBEDTLS_VERSION_MAJOR >= 3 #define CPPHTTPLIB_MBEDTLS_V3 #endif @@ -696,7 +708,7 @@ inline from_chars_result from_chars(const char *first, const char *last, return {first, std::errc::invalid_argument}; } - value = negative ? -result : result; + value = negative ? T(0) - result : result; return {p, std::errc{}}; } @@ -2957,7 +2969,18 @@ inline size_t get_header_value_u64(const Headers &headers, std::advance(it, static_cast(id)); if (it != rng.second) { if (is_numeric(it->second)) { - return static_cast(std::strtoull(it->second.data(), nullptr, 10)); + // Parse at size_t width so an out-of-range Content-Length is reported + // rather than silently saturated/truncated (a value above 2^32 would + // otherwise wrap to a small framing length on 32-bit builds). Flag it + // and return SIZE_MAX so the existing oversized-value guards reject it. + size_t val = 0; + const auto &s = it->second; + auto r = from_chars(s.data(), s.data() + s.size(), val); + if (r.ec == std::errc::result_out_of_range) { + is_invalid_value = true; + return (std::numeric_limits::max)(); + } + return val; } else { is_invalid_value = true; } @@ -3403,6 +3426,7 @@ bool is_obs_text(char c); bool is_field_vchar(char c); bool is_field_content(const std::string &s); bool is_field_value(const std::string &s); +bool is_field_valid(const std::string &name, const std::string &value); } // namespace fields } // namespace detail @@ -3422,8 +3446,11 @@ namespace impl { // setup callbacks (cast ctx_t to tls::impl::MbedTlsContext*). struct MbedTlsContext { mbedtls_ssl_config conf; +#ifndef CPPHTTPLIB_MBEDTLS_V4 + // Mbed TLS 4.x uses PSA Crypto's internal RNG; no explicit entropy/DRBG. mbedtls_entropy_context entropy; mbedtls_ctr_drbg_context ctr_drbg; +#endif mbedtls_x509_crt ca_chain; mbedtls_x509_crt own_cert; mbedtls_pk_context own_key; From 2cfc7670ed104a360b86beb55b6c06458015af7d Mon Sep 17 00:00:00 2001 From: Tekin Ertekin Date: Sat, 25 Jul 2026 21:53:08 +0300 Subject: [PATCH 5/9] server : add missing task parameters(adaptive_target, adaptive_decay) in generation_settings (#25830) These two parameters were overlooked when task parameters were being JSONized within `generation_settings` and have been added. A regression test has been added to prevent the problem from recurring and it passes. Fixes #25803 --- tools/server/server-task.cpp | 4 ++++ tools/server/tests/unit/test_completion.py | 2 ++ 2 files changed, 6 insertions(+) diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index f01780100c..1fd7cce27b 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -63,6 +63,8 @@ json task_params::to_json(bool only_metrics) const { {"mirostat", sampling.mirostat}, {"mirostat_tau", sampling.mirostat_tau}, {"mirostat_eta", sampling.mirostat_eta}, + {"adaptive_target", sampling.adaptive_target}, + {"adaptive_decay", sampling.adaptive_decay}, {"max_tokens", n_predict}, {"n_predict", n_predict}, // TODO: deduplicate? {"n_keep", n_keep}, @@ -114,6 +116,8 @@ json task_params::to_json(bool only_metrics) const { {"mirostat", sampling.mirostat}, {"mirostat_tau", sampling.mirostat_tau}, {"mirostat_eta", sampling.mirostat_eta}, + {"adaptive_target", sampling.adaptive_target}, + {"adaptive_decay", sampling.adaptive_decay}, {"stop", antiprompt}, {"max_tokens", n_predict}, {"n_predict", n_predict}, // TODO: deduplicate? diff --git a/tools/server/tests/unit/test_completion.py b/tools/server/tests/unit/test_completion.py index 1e0891987a..9375e0110e 100644 --- a/tools/server/tests/unit/test_completion.py +++ b/tools/server/tests/unit/test_completion.py @@ -66,6 +66,8 @@ def test_completion_stream(prompt: str, n_predict: int, re_content: str, n_promp assert server.n_predict is not None assert data["generation_settings"]["n_predict"] == min(n_predict, server.n_predict) assert data["generation_settings"]["seed"] == server.seed + assert "adaptive_target" in data["generation_settings"] + assert "adaptive_decay" in data["generation_settings"] assert match_regex(re_content, content) else: assert len(data["tokens"]) > 0 From abc348790edeb93ce1f9c6869c4fd6a07ca77d64 Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen Date: Sat, 25 Jul 2026 21:15:15 +0200 Subject: [PATCH 6/9] server: add format arg to datetime tool (#26117) --- tools/server/server-tools.cpp | 37 ++++++++++++++++++++++++++++++----- 1 file changed, 32 insertions(+), 5 deletions(-) diff --git a/tools/server/server-tools.cpp b/tools/server/server-tools.cpp index 9eb57abaea..a82a3d6025 100644 --- a/tools/server/server-tools.cpp +++ b/tools/server/server-tools.cpp @@ -7,6 +7,7 @@ #include #include #include +#include #include #include #include @@ -1048,16 +1049,42 @@ struct server_tool_get_datetime : server_tool { {"type", "function"}, {"function", { {"name", name}, - {"description", "Returns the current date and time"}, + {"description", "Returns the current date and time in UTC"}, + {"parameters", { + {"type", "object"}, + {"properties", { + {"format", { + {"type", "string"}, + {"description", + "strftime()-style format string for the output (default: \"%Y-%m-%dT%H:%M:%SZ\", " + "e.g. ISO 8601). Choose your own format if you need something else, " + "e.g. \"%A, %B %d %Y\" for a human-readable date."}, + }}, + }}, + }}, }}, }; } - json invoke(json, server_tool::stream *) const override { - auto now = std::chrono::system_clock::now(); - auto time = std::chrono::system_clock::to_time_t(now); + json invoke(json params, server_tool::stream *) const override { + std::string format = json_value(params, "format", std::string("%Y-%m-%dT%H:%M:%SZ")); - return {{"result", std::ctime(&time)}}; + auto now = std::chrono::system_clock::now(); + auto time = std::chrono::system_clock::to_time_t(now); + std::tm tm_utc; +#ifdef _WIN32 + gmtime_s(&tm_utc, &time); +#else + gmtime_r(&time, &tm_utc); +#endif + + char buf[256]; + size_t len = std::strftime(buf, sizeof(buf), format.c_str(), &tm_utc); + if (len == 0) { + return {{"error", "invalid format string"}}; + } + + return {{"result", std::string(buf, len)}}; } }; From c812c543f8ab480661bd10b9546515b608b4747f Mon Sep 17 00:00:00 2001 From: hogeheer499-commits Date: Sat, 25 Jul 2026 21:15:27 +0200 Subject: [PATCH 7/9] common : skip empty implicit default preset (#25643) The INI parser creates an implicit default section for top-level metadata. After reserved keys such as `version` are skipped, that section can have no model options but was still added and exposed in router mode. Skip only the empty implicit default while preserving real default presets, named presets, and the global `[*]` settings. Signed-off-by: JS van Dijk <267467744+hogeheer499-commits@users.noreply.github.com> Co-authored-by: JS van Dijk <267467744+hogeheer499-commits@users.noreply.github.com> --- common/preset.cpp | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/common/preset.cpp b/common/preset.cpp index 4362c0621b..eb0c60b09c 100644 --- a/common/preset.cpp +++ b/common/preset.cpp @@ -330,6 +330,10 @@ common_presets common_preset_context::load_from_ini(const std::string & path, co } } + if (preset.name == COMMON_PRESET_DEFAULT_NAME && preset.options.empty()) { + continue; + } + if (preset.name == "*") { // handle global preset global = preset; From 355303edab8745e722a2e0e1bd3f6a39399966f2 Mon Sep 17 00:00:00 2001 From: Todor Boinovski Date: Sat, 25 Jul 2026 15:47:29 -0700 Subject: [PATCH 8/9] hexagon: partial im2col support (#26007) * hexagon: add IM2COL op Add Hexagon IM2COL support targeting only patch-embedding convolutions. * hexagon: im2col refactor and cleanup * hex-im2col: instrument and update im2col. * hex-im2col: add local htp_vtcm_layout computation. --- ggml/src/ggml-hexagon/ggml-hexagon.cpp | 34 +++ ggml/src/ggml-hexagon/htp/CMakeLists.txt | 1 + ggml/src/ggml-hexagon/htp/htp-ctx.h | 1 + ggml/src/ggml-hexagon/htp/htp-ops.h | 1 + ggml/src/ggml-hexagon/htp/im2col-ops.c | 306 +++++++++++++++++++++++ ggml/src/ggml-hexagon/htp/main.c | 3 + 6 files changed, 346 insertions(+) create mode 100644 ggml/src/ggml-hexagon/htp/im2col-ops.c diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index 5db8fc84ba..bdb8af0820 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -3281,6 +3281,35 @@ static bool ggml_hexagon_supported_ssm_conv(const struct ggml_hexagon_session * GGML_UNUSED(sess); } +static bool ggml_hexagon_supported_im2col(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) { + const struct ggml_tensor * src1 = op->src[1]; + const struct ggml_tensor * dst = op; + + const bool is_2D = ((const int32_t *) op->op_params)[6] == 1; + if (!is_2D) { + return false; + } + + // For now support F32->F32 and F32->F16 only. + if (src1->type != GGML_TYPE_F32 || (dst->type != GGML_TYPE_F16 && dst->type != GGML_TYPE_F32)) { + return false; + } + + if (!ggml_is_contiguous(src1) || !ggml_is_contiguous(dst)) { + return false; + } + + // For now keep padded OPs on CPU. Will revisit once we expand coverage past patch-embed OPs. + const int32_t p0 = ((const int32_t *) op->op_params)[2]; + const int32_t p1 = ((const int32_t *) op->op_params)[3]; + if (p0 != 0 || p1 != 0) { + return false; + } + + GGML_UNUSED(sess); + return true; +} + static bool ggml_hexagon_supported_pad(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) { const struct ggml_tensor * src0 = op->src[0]; const struct ggml_tensor * dst = op; @@ -3430,6 +3459,7 @@ static htp_op_code op_remap_to_htp(const ggml_tensor * t) { case GGML_OP_SOLVE_TRI: return HTP_OP_SOLVE_TRI; case GGML_OP_TRI: return HTP_OP_TRI; case GGML_OP_PAD: return HTP_OP_PAD; + case GGML_OP_IM2COL: return HTP_OP_IM2COL; case GGML_OP_UNARY: switch (ggml_get_unary_op(t)) { @@ -4152,6 +4182,10 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons supp = ggml_hexagon_supported_ssm_conv(sess, op); break; + case GGML_OP_IM2COL: + supp = ggml_hexagon_supported_im2col(sess, op); + break; + case GGML_OP_GATED_DELTA_NET: supp = ggml_hexagon_supported_gated_delta_net(sess, op); break; diff --git a/ggml/src/ggml-hexagon/htp/CMakeLists.txt b/ggml/src/ggml-hexagon/htp/CMakeLists.txt index 4fb526f0c0..b00aa2bc94 100644 --- a/ggml/src/ggml-hexagon/htp/CMakeLists.txt +++ b/ggml/src/ggml-hexagon/htp/CMakeLists.txt @@ -42,6 +42,7 @@ add_library(${HTP_LIB} SHARED solve-tri-ops.c pad-ops.c argsort-ops.c + im2col-ops.c ) target_compile_definitions(${HTP_LIB} PRIVATE diff --git a/ggml/src/ggml-hexagon/htp/htp-ctx.h b/ggml/src/ggml-hexagon/htp/htp-ctx.h index 97b8c7f297..e0f9a0c40d 100644 --- a/ggml/src/ggml-hexagon/htp/htp-ctx.h +++ b/ggml/src/ggml-hexagon/htp/htp-ctx.h @@ -140,5 +140,6 @@ int op_diag(struct htp_ops_context * octx); int op_solve_tri(struct htp_ops_context * octx); int op_gated_delta_net(struct htp_ops_context * octx); int op_pad(struct htp_ops_context * octx); +int op_im2col(struct htp_ops_context * octx); #endif /* HTP_CTX_H */ diff --git a/ggml/src/ggml-hexagon/htp/htp-ops.h b/ggml/src/ggml-hexagon/htp/htp-ops.h index cad9a4f540..a138f062aa 100644 --- a/ggml/src/ggml-hexagon/htp/htp-ops.h +++ b/ggml/src/ggml-hexagon/htp/htp-ops.h @@ -98,6 +98,7 @@ enum htp_op_code { HTP_OP_NORM, HTP_OP_CONCAT, HTP_OP_CLAMP, + HTP_OP_IM2COL, HTP_OP_INVALID }; diff --git a/ggml/src/ggml-hexagon/htp/im2col-ops.c b/ggml/src/ggml-hexagon/htp/im2col-ops.c new file mode 100644 index 0000000000..35fc103df8 --- /dev/null +++ b/ggml/src/ggml-hexagon/htp/im2col-ops.c @@ -0,0 +1,306 @@ +#pragma clang diagnostic ignored "-Wunused-variable" +#pragma clang diagnostic ignored "-Wunused-function" +#pragma clang diagnostic ignored "-Wunused-but-set-variable" + +#include +#include +#include +#include +#include + +#define GGML_COMMON_DECL_C +#include "ggml-common.h" +#include "htp-ctx.h" +#include "htp-ops.h" +#include "hvx-utils.h" +#include "hex-dma.h" +#include "hex-profile.h" +#include "htp-vtcm.h" + +struct htp_im2col_context { + struct htp_ops_context * octx; + uint32_t npatches_per_thread; // patches = N*OH*OW (pure-DDR kernel) + + uint32_t pe_rows_per_thread; // N*OH rows per worker + uint32_t pe_src_row_bytes; // one output row's source: IC*KH*IW*4, rounded 256 + uint32_t pe_dst_row_bytes; // one output row's dst: OW*patch_stride*2, rounded 256 + + // Patch-embed DMA path VTCM ping-pong. + uint8_t * pe_vtcm_src; // base of the 2x src buffers region + uint8_t * pe_vtcm_dst; // base of the 2x dst buffers region + uint32_t pe_src_size_per_thread; // 2 * pe_src_row_bytes + uint32_t pe_dst_size_per_thread; // 2 * pe_dst_row_bytes +}; + +// Per-op VTCM layout for the patch-embed DMA path +struct htp_im2col_vtcm_layout { + size_t off_src; + size_t off_dst; + size_t src_bytes_per_thread; + size_t dst_bytes_per_thread; + size_t total_bytes; +}; + +static inline void htp_im2col_vtcm_layout_build(struct htp_im2col_vtcm_layout * L, + size_t src_row_bytes, + size_t dst_row_bytes, + uint32_t n_threads) { + L->src_bytes_per_thread = 2 * src_row_bytes; + L->dst_bytes_per_thread = 2 * dst_row_bytes; + + L->off_src = 0; + L->off_dst = L->off_src + L->src_bytes_per_thread * n_threads; + L->total_bytes = L->off_dst + L->dst_bytes_per_thread * n_threads; +} + +#define IM2COL_PATCHEMBED_BODY(FNAME, DST_CTYPE, COPY_FN, SPLAT_FN, DST_ELEM, TAG) \ + static void FNAME(unsigned int nth, unsigned int ith, void * data) { \ + struct htp_im2col_context * ictx = (struct htp_im2col_context *) data; \ + struct htp_ops_context * octx = ictx->octx; \ + struct htp_thread_trace * restrict tr = &octx->ctx->trace[ith]; \ + const struct htp_tensor * restrict src1 = octx->src[1]; \ + const struct htp_tensor * restrict dst = octx->dst; \ + const int32_t s0 = octx->op_params[0]; \ + const int32_t s1 = octx->op_params[1]; \ + const int32_t p0 = octx->op_params[2]; \ + const int32_t p1 = octx->op_params[3]; \ + const int32_t d0 = octx->op_params[4]; \ + const int32_t d1 = octx->op_params[5]; \ + const uint32_t N = src1->ne[3]; \ + const uint32_t IC = src1->ne[2]; \ + const uint32_t IH = src1->ne[1]; \ + const uint32_t IW = src1->ne[0]; \ + const uint32_t KH = octx->src[0]->ne[1]; \ + const uint32_t KW = octx->src[0]->ne[0]; \ + const uint32_t OH = dst->ne[2]; \ + const uint32_t OW = dst->ne[1]; \ + const uint32_t patch_stride = IC * KH * KW; \ + const float * restrict src_data = (const float *) src1->data; \ + DST_CTYPE * restrict dst_data = (DST_CTYPE *) dst->data; \ + const uint32_t npatches = N * OH * OW; \ + const uint32_t patch_start = ictx->npatches_per_thread * ith; \ + const uint32_t patch_end = MIN(patch_start + ictx->npatches_per_thread, npatches); \ + if (patch_start >= patch_end) { \ + return; \ + } \ + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, patch_start); \ + for (uint32_t p = patch_start; p < patch_end; p++) { \ + const uint32_t iow = p % OW; \ + const uint32_t ioh = (p / OW) % OH; \ + const uint32_t in = p / (OW * OH); \ + DST_CTYPE * restrict dst_patch = dst_data + (uint64_t) p * patch_stride; \ + for (uint32_t iic = 0; iic < IC; iic++) { \ + const float * restrict src_plane = src_data + ((uint64_t) in * IC + iic) * IH * IW; \ + for (uint32_t ikh = 0; ikh < KH; ikh++) { \ + const int32_t iih = (int32_t) ioh * s1 + (int32_t) ikh * d1 - p1; \ + DST_CTYPE * restrict out_run = dst_patch + iic * (KH * KW) + ikh * KW; \ + if (iih < 0 || iih >= (int32_t) IH) { \ + SPLAT_FN(out_run, 0.0f, KW); \ + continue; \ + } \ + const int32_t iiw0 = (int32_t) iow * s0 - p0; \ + const float * restrict src_run = src_plane + (uint64_t) iih * IW + iiw0; \ + if (d0 == 1) { \ + /* contiguous source run: [lo,hi) is in-bounds, tails are zero pad */ \ + const int32_t lo = iiw0 < 0 ? -iiw0 : 0; \ + int32_t hi = (int32_t) IW - iiw0; \ + if (hi > (int32_t) KW) { \ + hi = (int32_t) KW; \ + } \ + if (hi <= lo) { \ + SPLAT_FN(out_run, 0.0f, KW); \ + } else { \ + if (lo > 0) { \ + SPLAT_FN(out_run, 0.0f, (uint32_t) lo); \ + } \ + COPY_FN((uint8_t *) (out_run + lo), (const uint8_t *) (src_run + lo), \ + (uint32_t) (hi - lo)); \ + if (hi < (int32_t) KW) { \ + SPLAT_FN(out_run + hi, 0.0f, (KW - (uint32_t) hi)); \ + } \ + } \ + continue; \ + } \ + for (uint32_t ikw = 0; ikw < KW; ikw++) { \ + const int32_t iiw = (int32_t) iow * s0 + (int32_t) ikw * d0 - p0; \ + out_run[ikw] = (iiw < 0 || iiw >= (int32_t) IW) ? \ + (DST_CTYPE) 0.0f : \ + (DST_CTYPE) src_plane[(uint64_t) iih * IW + iiw]; \ + } \ + } \ + } \ + } \ + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, patch_start); \ + } + +IM2COL_PATCHEMBED_BODY(im2col_patchembed_thread, __fp16, hvx_copy_f16_f32_uu, hvx_splat_f16_u, sizeof(__fp16), "f32-f16") +IM2COL_PATCHEMBED_BODY(im2col_patchembed_f32_thread, float, hvx_copy_f32_uu, hvx_splat_f32_u, sizeof(float), "f32-f32") + +#define IM2COL_PATCHEMBED_DMA_BODY(FNAME, DST_CTYPE, COPY_FN, SPLAT_FN, DST_ELEM, TAG) \ + static void FNAME(unsigned int nth, unsigned int ith, void * data) { \ + struct htp_im2col_context * ictx = (struct htp_im2col_context *) data; \ + struct htp_ops_context * octx = ictx->octx; \ + struct htp_thread_trace * restrict tr = &octx->ctx->trace[ith]; \ + const struct htp_tensor * restrict src1 = octx->src[1]; \ + const struct htp_tensor * restrict dst = octx->dst; \ + const uint32_t N = src1->ne[3], IC = src1->ne[2], IH = src1->ne[1], IW = src1->ne[0]; \ + const uint32_t KH = octx->src[0]->ne[1], KW = octx->src[0]->ne[0]; \ + const uint32_t OH = dst->ne[2], OW = dst->ne[1]; \ + const uint32_t patch_stride = IC * KH * KW; \ + const float * restrict src_data = (const float *) src1->data; \ + DST_CTYPE * restrict dst_data = (DST_CTYPE *) dst->data; \ + dma_queue * dmaq = octx->ctx->dma[ith]; \ + uint8_t * src_base = ictx->pe_vtcm_src + ith * ictx->pe_src_size_per_thread; \ + uint8_t * dst_base = ictx->pe_vtcm_dst + ith * ictx->pe_dst_size_per_thread; \ + float * srcb = (float *) src_base; \ + DST_CTYPE * dstb = (DST_CTYPE *) dst_base; \ + const uint32_t nrows = N * OH; \ + const uint32_t per_thread = ictx->pe_rows_per_thread; \ + const uint32_t row_start = per_thread * ith; \ + const uint32_t row_end = MIN(row_start + per_thread, nrows); \ + if (row_start >= row_end) \ + return; \ + for (uint32_t r = row_start; r < row_end; r++) { \ + const uint32_t in = r / OH; \ + const uint32_t ioh = r % OH; \ + for (uint32_t ikh = 0; ikh < KH; ikh++) { \ + int32_t iih = (int32_t) ioh * (int32_t) KH + (int32_t) ikh; \ + int ok = (iih >= 0 && iih < (int32_t) IH); \ + for (uint32_t iic = 0; iic < IC; iic++) { \ + float * vdst = srcb + ((uint64_t) (iic * KH + ikh)) * IW; \ + const float * _vsrc = \ + ok ? (src_data + ((uint64_t) (in * IC + iic) * IH + iih) * IW) : (const float *) vdst; \ + dma_queue_push_ddr_to_vtcm( \ + dmaq, dma_make_ptr((uint8_t *) vdst, ok ? (const uint8_t *) _vsrc : (const uint8_t *) vdst), \ + IW * sizeof(float), IW * sizeof(float), ok ? 1 : 0); \ + } \ + } \ + for (uint32_t i = 0; i < IC * KH; i++) \ + dma_queue_pop(dmaq); \ + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, r); \ + for (uint32_t iow = 0; iow < OW; iow++) { \ + DST_CTYPE * dst_patch = dstb + (uint64_t) iow * patch_stride; \ + for (uint32_t ikh = 0; ikh < KH; ikh++) { \ + int32_t iih = (int32_t) ioh * (int32_t) KH + (int32_t) ikh; \ + for (uint32_t iic = 0; iic < IC; iic++) { \ + DST_CTYPE * out_run = dst_patch + iic * (KH * KW) + ikh * KW; \ + if (iih < 0 || iih >= (int32_t) IH) { \ + SPLAT_FN(out_run, 0.0f, KW); \ + continue; \ + } \ + const float * src_run = srcb + ((uint64_t) (iic * KH + ikh)) * IW + (uint64_t) iow * KW; \ + COPY_FN((uint8_t *) out_run, (const uint8_t *) src_run, KW); \ + } \ + } \ + } \ + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, r); \ + DST_CTYPE * ddr_row = dst_data + ((uint64_t) (in * OH + ioh) * OW) * patch_stride; \ + dma_queue_push_vtcm_to_ddr(dmaq, dma_make_ptr((uint8_t *) ddr_row, (uint8_t *) dstb), \ + OW * patch_stride * (DST_ELEM), OW * patch_stride * (DST_ELEM), 1); \ + dma_queue_flush(dmaq); \ + } \ + } + +IM2COL_PATCHEMBED_DMA_BODY(im2col_patchembed_dma_thread, __fp16, hvx_copy_f16_f32_uu, hvx_splat_f16_u, sizeof(__fp16), "pe-dma-f16") +IM2COL_PATCHEMBED_DMA_BODY(im2col_patchembed_dma_f32_thread, float, hvx_copy_f32_uu, hvx_splat_f32_u, sizeof(float), "pe-dma-f32") + +static bool im2col_use_patchembed_dma(const struct htp_ops_context * octx) { + const int32_t s0 = octx->op_params[0], s1 = octx->op_params[1]; + const int32_t p0 = octx->op_params[2], p1 = octx->op_params[3]; + const int32_t d0 = octx->op_params[4], d1 = octx->op_params[5]; + const int is_2D = octx->op_params[6] == 1; + if (!is_2D) { + return false; + } + if (octx->dst->type != HTP_TYPE_F16 && octx->dst->type != HTP_TYPE_F32) { + return false; + } + const uint32_t KH = octx->src[0]->ne[1], KW = octx->src[0]->ne[0]; + if (s0 != (int32_t) KW || s1 != (int32_t) KH) { + return false; // non-overlapping + } + if (p0 != 0 || p1 != 0) { + return false; // no padding + } + if (d0 != 1 || d1 != 1) { + return false; // no dilation + } + return true; +} + +// Sizes the per-thread 2x(src,dst) VTCM ping-pong for the patch-embed DMA path. +// Returns false if it doesn't fit the VTCM budget (caller falls back). +static bool im2col_patchembed_dma_fits(struct htp_ops_context * octx, + struct htp_im2col_context * ictx, + uint32_t n_threads) { + const uint32_t IC = octx->src[1]->ne[2], IW = octx->src[1]->ne[0]; + const uint32_t KH = octx->src[0]->ne[1], KW = octx->src[0]->ne[0]; + const uint32_t OW = octx->dst->ne[1]; + const uint32_t patch_stride = IC * KH * KW; + + ictx->pe_src_row_bytes = hex_round_up(IC * KH * IW * sizeof(float), 256); + const uint32_t dst_elem = (octx->dst->type == HTP_TYPE_F16) ? sizeof(__fp16) : sizeof(float); + ictx->pe_dst_row_bytes = hex_round_up(OW * patch_stride * dst_elem, 256); + + // 2 src + 2 dst buffers per thread (ping-pong), src region first then dst. + struct htp_im2col_vtcm_layout L; + htp_im2col_vtcm_layout_build(&L, ictx->pe_src_row_bytes, ictx->pe_dst_row_bytes, n_threads); + if (L.total_bytes > octx->ctx->vtcm_size) { + return false; + } + + uint8_t * const base = octx->ctx->vtcm_base; + ictx->pe_vtcm_src = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src); + ictx->pe_vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst); + ictx->pe_src_size_per_thread = (uint32_t) L.src_bytes_per_thread; + ictx->pe_dst_size_per_thread = (uint32_t) L.dst_bytes_per_thread; + return true; +} + +int op_im2col(struct htp_ops_context * octx) { + const struct htp_tensor * src1 = octx->src[1]; + const struct htp_tensor * dst = octx->dst; + + if (src1->type != HTP_TYPE_F32 || (dst->type != HTP_TYPE_F16 && dst->type != HTP_TYPE_F32)) { + FARF(ERROR, "im2col: only (F32 image -> F16/F32 columns) supported"); + return HTP_STATUS_NO_SUPPORT; + } + + const uint32_t N = src1->ne[3]; + const uint32_t OH = dst->ne[2]; + const uint32_t OW = dst->ne[1]; + const uint32_t npatches = N * OH * OW; + const uint32_t n_threads = MIN(octx->n_threads, npatches); + + if ((octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) || n_threads == 0) { + return HTP_STATUS_OK; + } + + struct htp_im2col_context ictx = { 0 }; + ictx.octx = octx; + ictx.npatches_per_thread = (npatches + n_threads - 1) / n_threads; + + // Clean non-overlapping patch-embed -> DMA kernel (if it fits VTCM); + // everything else (padding/dilation/stride edges) -> pure-DDR kernel. + if (im2col_use_patchembed_dma(octx)) { + const uint32_t nrows = N * OH; + const uint32_t pth = MIN(octx->n_threads, nrows); + if (pth > 0 && im2col_patchembed_dma_fits(octx, &ictx, pth)) { + ictx.pe_rows_per_thread = (nrows + pth - 1) / pth; + if (dst->type == HTP_TYPE_F16) { + work_queue_run(octx->ctx->work_queue, im2col_patchembed_dma_thread, &ictx, pth); + } else { + work_queue_run(octx->ctx->work_queue, im2col_patchembed_dma_f32_thread, &ictx, pth); + } + return HTP_STATUS_OK; + } + // else: doesn't fit -> fall through to the pure-DDR kernel below. + } + + if (dst->type == HTP_TYPE_F16) { + work_queue_run(octx->ctx->work_queue, im2col_patchembed_thread, &ictx, n_threads); + } else { + work_queue_run(octx->ctx->work_queue, im2col_patchembed_f32_thread, &ictx, n_threads); + } + return HTP_STATUS_OK; +} diff --git a/ggml/src/ggml-hexagon/htp/main.c b/ggml/src/ggml-hexagon/htp/main.c index 7d65e46436..880e20c995 100644 --- a/ggml/src/ggml-hexagon/htp/main.c +++ b/ggml/src/ggml-hexagon/htp/main.c @@ -781,6 +781,9 @@ static int execute_op(struct htp_ops_context * octx) { case HTP_OP_PAD: return op_pad(octx); + case HTP_OP_IM2COL: + return op_im2col(octx); + case HTP_OP_CONCAT: return op_concat(octx); From 20455a4ad336e958cfe8f82efce2c46cd44c4fa3 Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen Date: Sun, 26 Jul 2026 01:08:49 +0200 Subject: [PATCH 9/9] server: support MCP stdio (#26062) * move server_pipe to common * init impl * vendor: update subprocess.h * add server_mcp_stdio * stderr drain * server_mcp_transport * server_mcp_stdio is now framing-only, no json * internal/mcp-stdio: integration + tests + fixes (#26075) * server-mcp: harden transport and wire up the tool integration Builds on the transport/manager architecture (server_mcp_transport + server_pipe) with the hardening and integration the draft did not yet have. Hardening: * Reader and stderr pumps are polled (running-aware) instead of blocking on a read that only ends at EOF. subprocess_terminate() SIGKILLs only the direct child, so a grandchild the MCP server spawned that inherited the pipe would otherwise keep the write end open and hang teardown (both warmup shutdown at startup and process shutdown). The writer is likewise non-blocking + polled. * Windows: resolve the command through PATHEXT so "npx" (npm ships npx.cmd, never npx.exe) spawns, matching POSIX's PATH search; and enumerate the parent environment as UTF-8 (GetEnvironmentStringsW) instead of the active code page. * server_pipe gains an opt-in max_size (default unbounded, so the router's streaming use is unchanged); the MCP reply queue uses it so a server that streams unsolicited notifications between requests cannot grow it without bound. Integration: * --mcp-servers-config / --mcp-servers-json flags; enabling MCP restricts default CORS to localhost, same as --tools. * MCP tools are exposed through /tools (and chat-completions) as _, skipping names that collide with a built-in or another MCP tool. * Manager lifecycle wired into llama_server(): warmup at start, shutdown() from the signal handler before the HTTP server drains, blocking teardown in clean_up(). * SIGPIPE ignored so a child dying mid-write yields EPIPE rather than killing us. Assisted-By: Claude Opus 4.8 * server-mcp: add MCP test suite with grandchild deadlock regression test 21 tests over the /tools endpoint: tool discovery/invocation, timeouts, crash recovery and respawn cooldown, warmup partial failure, malformed and batched notification+response output, tool-definition shape, and prompt shutdown during a slow call. The last test spawns an MCP server that leaves a grandchild inheriting its stdout/stderr and asserts the server both starts and stops promptly. Verified it fails (5s SIGKILL fallback on a deadlocked reader-join) when the pump is made to ignore the running flag, and passes with the polled reader. Assisted-By: Claude Opus 4.8 * clean up * clean up 2 * even stricter life cycle * nits * nits 2 --------- Co-authored-by: Xuan Son Nguyen * fix some edge cases * fix last_error data race * fix response schema + docs * server: fix MCP zombie leak and timeout-induced transport teardown join_pumps() never reaped the child, leaking one zombie per spawn: call subprocess_join() before subprocess_destroy(). A per-call timeout permanently closed from_server and got a healthy transport evicted: add close_on_stop to server_pipe::read() and pass false from send_rpc(), where should_stop is a per-request deadline and a late reply is already skipped on id mismatch. Also drop the unreachable disconnect cancellation in server_mcp_tool::invoke(): support_stream is false, st is always null. (cherry picked from commit e6de1ec043174fd0570b1e60d47f06c7c19d620d) Assisted-by: Claude Opus 4.8 * server: make MCP test fixtures JSON-RPC 2.0 compliant Add the missing notification guard to mcp_malformed_server.py and mcp_burst_server.py (the latter treated id 0 as a notification and replied to unknown ones; its notification table is now unused). Return -32602 instead of -32601 for unknown tools: tools/call is a valid method, the tool name is the invalid parameter. Also fix the test module docstring: tools are named _. (cherry picked from commit 74a08e8c311dabf3b49d06cc6d754b0097ae7a38) Assisted-by: Claude Opus 4.8 --------- Co-authored-by: Piotr Wilkin (ilintar) Co-authored-by: Pascal --- common/arg.cpp | 21 +- common/common.h | 4 + tools/server/CMakeLists.txt | 2 + tools/server/README-dev.md | 4 +- tools/server/server-common.h | 72 +- tools/server/server-mcp.cpp | 836 ++++++++++++++++++ tools/server/server-mcp.h | 176 ++++ tools/server/server-models.cpp | 49 +- tools/server/server-tools.cpp | 73 +- tools/server/server-tools.h | 6 +- tools/server/server.cpp | 35 +- .../server/tests/fixtures/mcp_burst_server.py | 118 +++ .../server/tests/fixtures/mcp_crash_server.py | 114 +++ .../server/tests/fixtures/mcp_echo_server.py | 164 ++++ .../tests/fixtures/mcp_grandchild_server.py | 100 +++ .../tests/fixtures/mcp_malformed_server.py | 113 +++ .../server/tests/fixtures/mcp_slow_server.py | 132 +++ tools/server/tests/unit/test_mcp_servers.py | 718 +++++++++++++++ tools/server/tests/utils.py | 6 + 19 files changed, 2680 insertions(+), 63 deletions(-) create mode 100644 tools/server/server-mcp.cpp create mode 100644 tools/server/server-mcp.h create mode 100644 tools/server/tests/fixtures/mcp_burst_server.py create mode 100644 tools/server/tests/fixtures/mcp_crash_server.py create mode 100755 tools/server/tests/fixtures/mcp_echo_server.py create mode 100644 tools/server/tests/fixtures/mcp_grandchild_server.py create mode 100644 tools/server/tests/fixtures/mcp_malformed_server.py create mode 100644 tools/server/tests/fixtures/mcp_slow_server.py create mode 100644 tools/server/tests/unit/test_mcp_servers.py diff --git a/common/arg.cpp b/common/arg.cpp index a287b907d4..9753441313 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -850,8 +850,9 @@ static bool common_params_parse_ex(int argc, char ** argv, common_params_context params.kv_overrides.back().key[0] = 0; } - if (!params.server_tools.empty() && !params.cors_origins_explicit) { - LOG_WRN("server tools are enabled, using localhost as default CORS origin (change via --cors-origins)\n"); + const bool mcp_enabled = !params.mcp_servers_config.empty() || !params.mcp_servers_json.empty(); + if ((!params.server_tools.empty() || mcp_enabled) && !params.cors_origins_explicit) { + LOG_WRN("server tools or MCP servers are enabled, using localhost as default CORS origin (change via --cors-origins)\n"); params.cors_origins = "localhost"; } @@ -3261,6 +3262,22 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.server_tools = parse_csv_row(value); } ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_TOOLS")); + add_opt(common_arg( + {"--mcp-servers-config"}, "PATH", + "experimental: path to JSON file with MCP server definitions (Cursor-compatible format) - do not enable in untrusted environments (default: none)\n" + "note: for security reasons, this will limit --cors-origins to localhost by default", + [](common_params & params, const std::string & value) { + params.mcp_servers_config = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_MCP_SERVERS_CONFIG")); + add_opt(common_arg( + {"--mcp-servers-json"}, "JSON", + "experimental: inline JSON with MCP server definitions (Cursor-compatible format) - do not enable in untrusted environments (default: none)\n" + "note: for security reasons, this will limit --cors-origins to localhost by default", + [](common_params & params, const std::string & value) { + params.mcp_servers_json = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_MCP_SERVERS_JSON")); add_opt(common_arg( {"-ag", "--agent"}, {"-no-ag", "--no-agent"}, diff --git a/common/common.h b/common/common.h index b5687c1083..2792521836 100644 --- a/common/common.h +++ b/common/common.h @@ -668,6 +668,10 @@ struct common_params { // enable built-in tools std::vector server_tools; + // MCP server configs (Cursor-compatible JSON) + std::string mcp_servers_config; // path to JSON file with MCP server definitions + std::string mcp_servers_json; // inline JSON with MCP server definitions + // router server configs std::string models_dir = ""; // directory containing models for the router server std::string models_preset = ""; // directory containing model presets for the router server diff --git a/tools/server/CMakeLists.txt b/tools/server/CMakeLists.txt index b5c40884fd..280bd9e19d 100644 --- a/tools/server/CMakeLists.txt +++ b/tools/server/CMakeLists.txt @@ -19,6 +19,8 @@ add_library(${TARGET} STATIC server-stream.h server-tools.cpp server-tools.h + server-mcp.cpp + server-mcp.h server-schema.cpp server-schema.h ) diff --git a/tools/server/README-dev.md b/tools/server/README-dev.md index e81336e5e2..b4ec9f17d3 100644 --- a/tools/server/README-dev.md +++ b/tools/server/README-dev.md @@ -189,7 +189,7 @@ This endpoint is intended to be used internally by the Web UI and subject to cha Get a list of tools, each tool has these fields: - `tool` (string): the ID name of the tool, to be used in POST call. Example: `read_file` - `display_name` (string): the name to be displayed on UI. Example: `Read file` -- `type` (string): always be `"builtin"` for now +- `type` (string): `"builtin"` for a built-in tool, or `"mcp"` for a tool exposed by an MCP server - `permissions` (object): a mapping string --> boolean that indicates the permission required by this tool. This is useful for the UI to ask the user before calling the tool. For now, the only permission supported is `"write"` - `definition` (object): the OAI-compat definition of this tool @@ -199,7 +199,7 @@ Invoke a tool call, request body is a JSON object with: - `tool` (string): the name of the tool - `params` (object): a mapping from argument name (string) to argument value -Returns JSON object. There are two response formats: +Returns JSON object. There are two response formats (MCP tools use the same two formats: their result content is concatenated into `plain_text_response`, and RPC or tool errors are surfaced as the `error` string): Format 1: Plain text. The text will be placed into a field called `plain_text_response`, example: diff --git a/tools/server/server-common.h b/tools/server/server-common.h index 5837366380..6ef797ebb4 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -9,9 +9,15 @@ #define JSON_ASSERT GGML_ASSERT #include +#include +#include +#include +#include +#include +#include +#include #include #include -#include using json = nlohmann::ordered_json; @@ -376,3 +382,67 @@ server_tokens format_prompt_rerank( mtmd_context * mctx, const std::string & query, const std::string & doc); + +// simple implementation of a pipe +// used for streaming data between threads +template +struct server_pipe { + std::mutex mutex; + std::condition_variable cv; + std::queue queue; + std::atomic writer_closed{false}; + std::atomic reader_closed{false}; + + // 0 = unbounded (default) + // > 0, write() drops the oldest item once the queue is full + size_t max_size = 0; + + void close_write() { + writer_closed.store(true, std::memory_order_relaxed); + cv.notify_all(); + } + + void close_read() { + reader_closed.store(true, std::memory_order_relaxed); + cv.notify_all(); + } + + // close_on_stop = true: should_stop means the reader is gone for good, so the writer is told the pipe is broken. + // close_on_stop = false: should_stop is a per-read deadline and further reads still come, so the pipe stays usable. + bool read(T & output, const std::function & should_stop, bool close_on_stop = true) { + std::unique_lock lk(mutex); + constexpr auto poll_interval = std::chrono::milliseconds(500); + while (true) { + if (!queue.empty()) { + output = std::move(queue.front()); + queue.pop(); + return true; + } + if (writer_closed.load()) { + return false; // clean EOF + } + if (should_stop && should_stop()) { // a null should_stop means "never stop" + if (close_on_stop) { + close_read(); // signal broken pipe to writer + } + return false; // cancelled / deadline reached + } + cv.wait_for(lk, poll_interval); + } + } + + bool write(T && data) { + std::lock_guard lk(mutex); + if (reader_closed.load()) { + return false; // broken pipe + } + if (max_size > 0) { + while (queue.size() >= max_size) { + queue.pop(); // drop oldest to stay bounded + } + } + queue.push(std::move(data)); + cv.notify_one(); + return true; + } +}; diff --git a/tools/server/server-mcp.cpp b/tools/server/server-mcp.cpp new file mode 100644 index 0000000000..aa87afeb06 --- /dev/null +++ b/tools/server/server-mcp.cpp @@ -0,0 +1,836 @@ +#include "server-mcp.h" + +#include + +#include +#include +#include +#include +#include +#include +#include + +#if defined(_WIN32) +# include +# include +#else +# include +# include +# include +# include +extern char ** environ; +#endif + +// read NDJSON lines from a child pipe, calling on_line per line until `running` clears, EOF/error, or on_line returns false. +// polled, not blocking: a grandchild can inherit the pipe's write end and hold it open (terminate() kills only the direct child), so a blocking read would hang teardown on an EOF that never comes. +static void mcp_pump_ndjson(FILE * f, std::atomic & running, + const std::function & on_line) { + if (!f) { + return; + } + const int poll_ms = 50; + const size_t max_line = 8 * 1024 * 1024; // drop any single NDJSON line larger than this, so a child that never emits '\n' can't grow buf without bound +#if defined(_WIN32) + HANDLE h = (HANDLE) _get_osfhandle(_fileno(f)); +#else + int fd = fileno(f); + int fl = fcntl(fd, F_GETFL, 0); + if (fl >= 0) { + fcntl(fd, F_SETFL, fl | O_NONBLOCK); + } +#endif + std::string buf; + bool skipping = false; // discarding an over-long line until its terminating newline + char chunk[4096]; + while (running.load()) { + size_t n = 0; +#if defined(_WIN32) + DWORD avail = 0; + if (!PeekNamedPipe(h, NULL, 0, NULL, &avail, NULL)) { + break; // pipe broken / child gone + } + if (avail == 0) { + std::this_thread::sleep_for(std::chrono::milliseconds(poll_ms)); + continue; + } + DWORD to_read = avail < (DWORD) sizeof(chunk) ? avail : (DWORD) sizeof(chunk); + DWORD got = 0; + if (!ReadFile(h, chunk, to_read, &got, NULL) || got == 0) { + break; + } + n = (size_t) got; +#else + struct pollfd pfd; + pfd.fd = fd; + pfd.events = POLLIN; + pfd.revents = 0; + int pr = poll(&pfd, 1, poll_ms); + if (pr < 0) { + if (errno == EINTR) { + continue; + } + break; + } + if (pr == 0) { + continue; // timeout -> re-check running + } + if (pfd.revents & (POLLERR | POLLNVAL)) { + break; + } + ssize_t r = read(fd, chunk, sizeof(chunk)); + if (r < 0) { + if (errno == EINTR || errno == EAGAIN || errno == EWOULDBLOCK) { + continue; + } + break; + } + if (r == 0) { + break; // EOF: child (and any pipe writers) closed the stream + } + n = (size_t) r; +#endif + buf.append(chunk, n); + + // resync after an over-long, unterminated line: discard bytes until the next newline + if (skipping) { + size_t nl = buf.find('\n'); + if (nl == std::string::npos) { + if (buf.size() > max_line) { + buf.clear(); // stay bounded while waiting for a terminator + } + continue; + } + buf.erase(0, nl + 1); + skipping = false; + } + + size_t pos; + while ((pos = buf.find('\n')) != std::string::npos) { + std::string line = buf.substr(0, pos); + buf.erase(0, pos + 1); + if (!line.empty() && line.back() == '\r') { + line.pop_back(); + } + if (line.empty()) { + continue; + } + if (!on_line(std::move(line))) { + return; + } + } + + // a partial line already larger than the cap and still no newline: drop it to avoid unbounded growth + if (buf.size() > max_line) { + SRV_WRN("MCP: dropping oversized line (> %zu bytes) from child pipe\n", max_line); + buf.clear(); + skipping = true; + } + } +} + +// +// server_mcp_server_config +// + +std::vector server_mcp_server_config::parse_from_json(const std::string & json_str) { + return parse_cursor_format(json::parse(json_str)); +} + +std::vector server_mcp_server_config::parse_cursor_format(const json & j) { + std::vector result; + + if (!j.contains("mcpServers") || !j.at("mcpServers").is_object()) { + return result; + } + + for (const auto & [name, cfg] : j.at("mcpServers").items()) { + server_mcp_server_config sc; + sc.name = name; + sc.command = cfg.value("command", std::string()); + sc.cwd = cfg.value("cwd", std::string()); + sc.timeout_ms = cfg.value("timeout_ms", sc.timeout_ms); + + if (cfg.contains("args") && cfg.at("args").is_array()) { + for (const auto & a : cfg.at("args")) { + sc.args.push_back(a.get()); + } + } + if (cfg.contains("env") && cfg.at("env").is_object()) { + for (const auto & [k, v] : cfg.at("env").items()) { + sc.env[k] = v.get(); + } + } + + if (sc.command.empty()) { + SRV_WRN("MCP server '%s' has no command, skipping\n", name.c_str()); + continue; + } + result.push_back(std::move(sc)); + } + + return result; +} + + +// +// server_mcp_transport +// + +static constexpr const char * MCP_PROTOCOL_VERSION = "2024-11-05"; + +static std::string rpc_error_message(const json & resp) { + if (resp.contains("error")) { + const json & e = resp.at("error"); + if (e.is_object()) { + return e.value("message", "unknown error"); + } + if (e.is_string()) { + return e.get(); + } + } + return "unknown error"; +} + +// normalize an MCP tools/call result to the /tools contract (see README-dev.md): +// concat text parts of result.content[], and surface an isError result +static json mcp_result_to_response(const json & result) { + std::string text; + if (result.contains("content") && result.at("content").is_array()) { + for (const auto & part : result.at("content")) { + if (part.is_object() && part.value("type", "") == "text") { + if (!text.empty()) { + text += "\n"; + } + text += part.value("text", ""); + } + } + } + if (result.is_object() && result.value("isError", false)) { + return {{"error", text.empty() ? "MCP tool returned an error" : text}}; + } + return {{"plain_text_response", text}}; +} + +json server_mcp_transport::send_rpc(const json & request, const std::function & should_stop) { + if (!to_server.write(request.dump())) { + return {{"error", {{"code", -32603}, {"message", "transport closed"}}}}; + } + + const bool has_id = request.contains("id"); + const auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(timeout_ms); + auto stop = [&]() { + return (should_stop && should_stop()) || std::chrono::steady_clock::now() >= deadline; + }; + + std::string frame; + while (from_server.read(frame, stop, false)) { + json reply; + try { + reply = json::parse(frame); + } catch (...) { + if (std::chrono::steady_clock::now() >= deadline) { + break; + } + continue; // skip malformed frame + } + // no id: a notification. mismatched id: a stale reply from a timed-out request (ids are monotonic, never a future one) + if (!has_id || (reply.contains("id") && reply.at("id") == request.at("id"))) { + return reply; + } + if (std::chrono::steady_clock::now() >= deadline) { + break; // a flood of notifications must not outrun the deadline + } + } + + if (should_stop && should_stop()) { + return {{"error", {{"code", -32603}, {"message", "cancelled"}}}}; + } + if (std::chrono::steady_clock::now() >= deadline) { + return {{"error", {{"code", -32603}, {"message", "request timed out"}}}}; + } + return {{"error", {{"code", -32603}, {"message", "transport closed"}}}}; +} + +bool server_mcp_transport::ensure_init(const std::function & should_stop) { + if (initialized) { + return true; + } + + json init_req = { + {"jsonrpc", "2.0"}, + {"id", next_id++}, + {"method", "initialize"}, + {"params", { + {"protocolVersion", MCP_PROTOCOL_VERSION}, + {"capabilities", json::object()}, + {"clientInfo", {{"name", "llama.cpp"}, {"version", "1.0"}}}, + }}, + }; + json resp = send_rpc(init_req, should_stop); + if (!resp.contains("result")) { + last_error = "initialize failed: " + rpc_error_message(resp); + return false; + } + + // notifications/initialized: no id, no reply expected + json notif = {{"jsonrpc", "2.0"}, {"method", "notifications/initialized"}}; + to_server.write(notif.dump()); + + initialized = true; + return true; +} + +std::vector server_mcp_transport::list_tools(const std::function & should_stop) { + std::lock_guard lock(rpc_mutex); + if (!ensure_init(should_stop)) { + return {}; + } + if (!tools.empty()) { + return tools; + } + + json req = {{"jsonrpc", "2.0"}, {"id", next_id++}, {"method", "tools/list"}}; + json resp = send_rpc(req, should_stop); + if (!resp.contains("result")) { + last_error = "tools/list failed: " + rpc_error_message(resp); + return {}; + } + + const json & result = resp.at("result"); + if (result.contains("tools") && result.at("tools").is_array()) { + for (const auto & t : result.at("tools")) { + server_mcp_tool_def def; + def.server_name = name; + def.name = t.value("name", ""); + def.description = t.value("description", ""); + if (t.contains("inputSchema")) { + def.input_schema = t.at("inputSchema"); + } + tools.push_back(std::move(def)); + } + } + return tools; +} + +json server_mcp_transport::call_tool(const std::string & tool_name, + const json & arguments, + const std::function & should_stop) { + std::lock_guard lock(rpc_mutex); + if (!ensure_init(should_stop)) { + return {{"error", last_error}}; + } + + json req = { + {"jsonrpc", "2.0"}, + {"id", next_id++}, + {"method", "tools/call"}, + {"params", {{"name", tool_name}, {"arguments", arguments}}}, + }; + json resp = send_rpc(req, should_stop); + if (resp.contains("error")) { + return {{"error", rpc_error_message(resp)}}; + } + if (resp.contains("result")) { + return mcp_result_to_response(resp.at("result")); + } + return {{"error", "invalid response from MCP server"}}; +} + +// +// server_mcp_stdio +// + +struct server_mcp_stdio::process_handle { + subprocess_s sp; + FILE * in = nullptr; // child stdin + FILE * out = nullptr; // child stdout + FILE * err = nullptr; // child stderr +}; + +#if defined(_WIN32) +// config strings are UTF-8 (from JSON) and subprocess.h converts them with CP_UTF8, so inputs must be UTF-8, not the active code page +static std::wstring windows_utf8_to_wide(const std::string & s) { + if (s.empty()) { + return std::wstring(); + } + int n = MultiByteToWideChar(CP_UTF8, 0, s.data(), (int) s.size(), NULL, 0); + if (n <= 0) { + return std::wstring(); + } + std::wstring w((size_t) n, L'\0'); + MultiByteToWideChar(CP_UTF8, 0, s.data(), (int) s.size(), &w[0], n); + return w; +} + +static std::string windows_wide_to_utf8(const wchar_t * s, int len /* -1 for NUL-terminated */) { + int n = WideCharToMultiByte(CP_UTF8, 0, s, len, NULL, 0, NULL, NULL); + if (n <= 0) { + return std::string(); + } + std::string out((size_t) n, '\0'); + WideCharToMultiByte(CP_UTF8, 0, s, len, &out[0], n, NULL, NULL); + if (len == -1 && !out.empty() && out.back() == '\0') { + out.pop_back(); // drop the terminator WideCharToMultiByte counts for -1 + } + return out; +} +#endif + +static std::string mcp_resolve_command(const std::string & command) { +#if defined(_WIN32) + // For Windows: make sure we handle ".exe" correctly, as well as UTF-8 + std::wstring wcmd = windows_utf8_to_wide(command); + wchar_t buf[MAX_PATH * 4]; + const DWORD cap = (DWORD) (sizeof(buf) / sizeof(buf[0])); + + auto search = [&](const wchar_t * ext) -> std::string { + DWORD n = SearchPathW(NULL, wcmd.c_str(), ext, cap, buf, NULL); + return (n > 0 && n < cap) ? windows_wide_to_utf8(buf, (int) n) : std::string(); + }; + + std::string found = search(NULL); // exact path / already-extensioned / .exe on PATH + if (!found.empty()) { + return found; + } + + std::wstring pathext; + DWORD need = GetEnvironmentVariableW(L"PATHEXT", NULL, 0); + if (need > 0) { + pathext.resize(need); + DWORD got = GetEnvironmentVariableW(L"PATHEXT", &pathext[0], need); + pathext.resize(got); + } + if (pathext.empty()) { + pathext = L".COM;.EXE;.BAT;.CMD"; + } + for (size_t start = 0; start <= pathext.size();) { + size_t sep = pathext.find(L';', start); + std::wstring ext = pathext.substr(start, sep == std::wstring::npos ? std::wstring::npos : sep - start); + if (!ext.empty()) { + found = search(ext.c_str()); + if (!found.empty()) { + return found; + } + } + if (sep == std::wstring::npos) { + break; + } + start = sep + 1; + } + return command; // give up and let subprocess.h report the spawn error +#else + return command; +#endif // _WIN32 +} + +static std::vector mcp_parent_env() { + std::vector env; +#if defined(_WIN32) + LPWCH block = GetEnvironmentStringsW(); + if (block) { + for (LPWCH e = block; *e; e += wcslen(e) + 1) { + env.emplace_back(windows_wide_to_utf8(e, -1)); + } + FreeEnvironmentStringsW(block); + } +#else + if (environ) { + for (char ** e = environ; *e; ++e) { + env.emplace_back(*e); + } + } +#endif + return env; +} + +// parent env with the config overrides applied, in "KEY=VALUE" form +static std::vector mcp_build_env(const std::map & overrides) { + std::vector env; + for (auto & e : mcp_parent_env()) { + size_t eq = e.find('='); + std::string key = eq == std::string::npos ? e : e.substr(0, eq); + if (overrides.find(key) == overrides.end()) { + env.push_back(e); + } + } + for (auto & [k, v] : overrides) { + env.push_back(k + "=" + v); + } + return env; +} + +server_mcp_stdio::server_mcp_stdio(const server_mcp_server_config & config) : config(config) { + name = config.name; + timeout_ms = config.timeout_ms; + // bound the reply queue: send_rpc only drains during a call, so unsolicited notifications would otherwise grow it without limit + from_server.max_size = 65536; +} + +server_mcp_stdio::~server_mcp_stdio() { + join_pumps(); +} + +bool server_mcp_stdio::start() { + std::vector argv_s; + argv_s.push_back(mcp_resolve_command(config.command)); + argv_s.insert(argv_s.end(), config.args.begin(), config.args.end()); + + int options = subprocess_option_no_window | subprocess_option_search_user_path; + std::vector envp_s; + if (config.env.empty()) { + options |= subprocess_option_inherit_environment; + } else { + envp_s = mcp_build_env(config.env); + } + + auto to_ptrs = [](std::vector & v) { + std::vector p; + p.reserve(v.size() + 1); + for (auto & s : v) { + p.push_back(s.c_str()); + } + p.push_back(nullptr); + return p; + }; + auto argv = to_ptrs(argv_s); + auto envp = to_ptrs(envp_s); + + auto handle = std::make_unique(); + int rc = subprocess_create_ex(argv.data(), options, + config.env.empty() ? nullptr : envp.data(), + config.cwd.empty() ? nullptr : config.cwd.c_str(), + &handle->sp); + if (rc != 0) { + SRV_WRN("MCP '%s': failed to spawn '%s'\n", config.name.c_str(), config.command.c_str()); + return false; + } + handle->in = subprocess_stdin(&handle->sp); + handle->out = subprocess_stdout(&handle->sp); + handle->err = subprocess_stderr(&handle->sp); + + proc = std::move(handle); + running.store(true); + reader = std::thread([this] { reader_loop(); }); + writer = std::thread([this] { writer_loop(); }); + errlog = std::thread([this] { errlog_loop(); }); + return true; +} + +void server_mcp_stdio::close() { + join_pumps(); +} + +bool server_mcp_stdio::is_alive() const { + return running.load(); +} + +std::string server_mcp_stdio::diagnostics() { + std::string out; + { + std::lock_guard lock(rpc_mutex); // last_error is written by send_rpc's callers + out = last_error; + } + std::lock_guard lk(err_mu); + if (!err_tail.empty()) { + if (!out.empty()) { + out += "; "; + } + out += "last stderr: " + err_tail; + } + return out; +} + +void server_mcp_stdio::reader_loop() { + mcp_pump_ndjson(proc->out, running, [this](std::string && line) { + return from_server.write(std::move(line)); // false => consumer gone, stop + }); + running.store(false); + to_server.close_write(); // stop the writer + from_server.close_write(); // EOF to any waiting caller +} + +// write all of `data` to child stdin, non-blocking and polled so teardown never hangs (a grandchild can hold the read end of a full pipe open). returns false on error/close/shutdown. +static bool mcp_write_all(FILE * f, const std::string & data, std::atomic & running) { + if (!f) { + return false; + } + size_t total = 0; +#if defined(_WIN32) + HANDLE h = (HANDLE) _get_osfhandle(_fileno(f)); + DWORD nowait = PIPE_NOWAIT; + SetNamedPipeHandleState(h, &nowait, NULL, NULL); + while (total < data.size() && running.load()) { + DWORD written = 0; + BOOL ok = WriteFile(h, data.data() + total, (DWORD) (data.size() - total), &written, NULL); + if (ok && written > 0) { + total += written; + continue; + } + if (!ok) { + DWORD err = GetLastError(); + if (err != ERROR_NO_DATA && err != ERROR_PIPE_BUSY) { + return false; + } + } + // backpressure (pipe full) is rare for small JSON-RPC frames; sleep rather than spin. + // no writable-wait exists for a PIPE_NOWAIT anonymous pipe, so this polls like the POSIX poll() path. + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } +#else + int fd = fileno(f); + int fl = fcntl(fd, F_GETFL, 0); + if (fl >= 0) { + fcntl(fd, F_SETFL, fl | O_NONBLOCK); + } + while (total < data.size() && running.load()) { + ssize_t n = write(fd, data.data() + total, data.size() - total); + if (n > 0) { + total += (size_t) n; + continue; + } + if (n == 0) { + return false; + } + if (errno == EINTR) { + continue; + } + if (errno != EAGAIN && errno != EWOULDBLOCK) { + return false; + } + struct pollfd pfd; + pfd.fd = fd; + pfd.events = POLLOUT; + pfd.revents = 0; + int pr = poll(&pfd, 1, 50); + if (pr < 0) { + if (errno == EINTR) { + continue; + } + return false; + } + if (pfd.revents & (POLLERR | POLLNVAL | POLLHUP)) { + return false; + } + } +#endif + return total == data.size(); +} + +void server_mcp_stdio::writer_loop() { + auto should_stop = [this] { return !running.load(); }; + std::string msg; + while (to_server.read(msg, should_stop)) { + msg.push_back('\n'); + if (!mcp_write_all(proc->in, msg, running)) { + break; // child gone or shutting down + } + } + running.store(false); + to_server.close_read(); // fail fast on any further send_rpc write + from_server.close_write(); // wake any caller waiting for a reply +} + +void server_mcp_stdio::errlog_loop() { + static constexpr size_t ERR_TAIL_MAX = 4096; + // drain stderr (an undrained pipe blocks the child): + // log it, and keep a bounded tail for reporting when the server dies + mcp_pump_ndjson(proc->err, running, [this](std::string && line) { + SRV_DBG("MCP '%s' stderr: %s\n", name.c_str(), line.c_str()); + std::lock_guard lk(err_mu); + err_tail += line; + err_tail += '\n'; + if (err_tail.size() > ERR_TAIL_MAX) { + err_tail.erase(0, err_tail.size() - ERR_TAIL_MAX); + } + return true; + }); +} + +void server_mcp_stdio::join_pumps() { + if (!proc) { + return; + } + running.store(false); + to_server.close_write(); // wake the writer if it waits for a message + from_server.close_write(); // wake any caller waiting for a reply + + subprocess_terminate(&proc->sp); // child death unblocks the blocked fread/fwrite + + if (writer.joinable()) writer.join(); + if (reader.joinable()) reader.join(); + if (errlog.joinable()) errlog.join(); + + subprocess_join(&proc->sp, nullptr); // reap the child: destroy() never waits, so the pid would stay a zombie for the process lifetime + subprocess_destroy(&proc->sp); // safe now: no thread touches the FILE* anymore + proc.reset(); +} + + +// +// server_mcp +// + +static constexpr int MCP_COOLDOWN_SECONDS = 5; +static constexpr int MCP_WARMUP_TIMEOUT_SECONDS = 10; // cap per-server tool discovery at startup + +server_mcp::~server_mcp() { + shutdown(); + + std::vector> to_close; + { + std::lock_guard lock(mutex); + for (auto & [name, t] : transports) { + to_close.push_back(std::move(t)); + } + transports.clear(); + } + for (auto & t : to_close) { + t->close(); + } +} + +std::shared_ptr server_mcp::create_transport(const server_mcp_server_config & cfg) { + return std::make_shared(cfg); +} + +void server_mcp::shutdown() { + stopping.store(true); +} + +const server_mcp_server_config * server_mcp::find_config(const std::string & name) const { + for (const auto & c : configs) { + if (c.name == name) { + return &c; + } + } + return nullptr; +} + +void server_mcp::start(const common_params & params) { + auto append = [this](const std::string & json_str) { + try { + auto parsed = server_mcp_server_config::parse_from_json(json_str); + if (parsed.empty()) { + SRV_WRN("%s", "MCP config: no servers found in JSON\n"); + } + for (auto & p : parsed) { + // names must be unique across both config sources: get_or_create / find_config key on the name + if (find_config(p.name)) { + SRV_WRN("MCP config: duplicate server name '%s', skipping\n", p.name.c_str()); + continue; + } + configs.push_back(std::move(p)); + } + } catch (const std::exception & e) { + throw std::runtime_error(std::string("failed to parse MCP config JSON: ") + e.what()); + } + }; + if (!params.mcp_servers_config.empty()) { + std::ifstream f = fs_open_ifstream(params.mcp_servers_config, std::ios::in); + if (!f) { + throw std::runtime_error("failed to open MCP config file: " + params.mcp_servers_config); + } + std::stringstream ss; + ss << f.rdbuf(); + append(ss.str()); + } + if (!params.mcp_servers_json.empty()) { + append(params.mcp_servers_json); + } + + if (configs.empty()) { + return; + } + + std::vector discovered; + for (const auto & cfg : configs) { + auto t = create_transport(cfg); + if (!t->start()) { + SRV_WRN("MCP warmup: failed to spawn '%s': %s\n", cfg.name.c_str(), t->diagnostics().c_str()); + continue; + } + // bound warmup per server so an unresponsive one can't stall startup for the full per-call timeout + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(MCP_WARMUP_TIMEOUT_SECONDS); + auto should_stop = [this, deadline]() { + return stopping.load() || std::chrono::steady_clock::now() >= deadline; + }; + auto tools = t->list_tools(should_stop); + SRV_INF("MCP warmup: '%s' discovered %zu tools\n", cfg.name.c_str(), tools.size()); + discovered.insert(discovered.end(), tools.begin(), tools.end()); + t->close(); + } + + std::lock_guard lock(mutex); + registry.swap(discovered); +} + +std::vector server_mcp::list_tools() const { + std::lock_guard lock(mutex); + return registry; +} + +json server_mcp::call_tool(const std::string & server_name, + const std::string & tool_name, + const json & arguments, + const std::function & should_stop) { + auto transport = get_or_create(server_name); + if (!transport) { + return {{"error", "MCP server unavailable: " + server_name}}; + } + + auto stop = [this, &should_stop]() { + return stopping.load() || (should_stop && should_stop()); + }; + return transport->call_tool(tool_name, arguments, stop); +} + +std::shared_ptr server_mcp::get_or_create(const std::string & name) { + std::vector> to_close; // closed after unlock + std::shared_ptr result; + + { + std::lock_guard lock(mutex); + if (stopping.load()) { + return nullptr; + } + + auto now = std::chrono::steady_clock::now(); + auto dead_it = dead_servers.find(name); + if (dead_it != dead_servers.end()) { + if (now < dead_it->second) { + return nullptr; + } + dead_servers.erase(dead_it); + } + + auto it = transports.find(name); + if (it != transports.end()) { + if (it->second->is_alive()) { + return it->second; + } + SRV_WRN("MCP '%s' is no longer alive: %s\n", name.c_str(), it->second->diagnostics().c_str()); + to_close.push_back(std::move(it->second)); + transports.erase(it); + } + + const server_mcp_server_config * cfg = find_config(name); + if (cfg) { + auto fresh = create_transport(*cfg); + if (fresh->start() && fresh->is_alive()) { + transports[name] = fresh; + result = fresh; + } else { + SRV_WRN("MCP '%s': failed to start: %s\n", name.c_str(), fresh->diagnostics().c_str()); + to_close.push_back(std::move(fresh)); + dead_servers[name] = now + std::chrono::seconds(MCP_COOLDOWN_SECONDS); + } + } + } + + for (auto & t : to_close) { + t->close(); // blocking call, no leaks + } + + return result; +} + diff --git a/tools/server/server-mcp.h b/tools/server/server-mcp.h new file mode 100644 index 0000000000..c7f33a3f79 --- /dev/null +++ b/tools/server/server-mcp.h @@ -0,0 +1,176 @@ +#pragma once + +#include "server-common.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// +// Configuration (Cursor-compatible "mcpServers" JSON) +// + +struct server_mcp_server_config { + std::string name; // config key, e.g. "filesystem" + std::string command; + std::vector args; + std::map env; // merged over the parent env + std::string cwd; + int timeout_ms = 30000; // per-tool-call timeout + + // throw on parse errors; missing "mcpServers" yields an empty list; entries without a "command" are skipped + static std::vector parse_from_json(const std::string & json_str); + static std::vector parse_cursor_format(const json & j); +}; + +// a tool advertised by an MCP server +struct server_mcp_tool_def { + std::string server_name; + std::string name; // bare tool name, no "_" prefix + std::string description; + json input_schema; // JSON Schema for the arguments, or null +}; + +// +// server_mcp_transport: one MCP server session. +// +// caller --send_rpc--> to_server --[writer]--> framing --> server +// caller <--send_rpc-- from_server <--[reader]-- framing <-- server +// +// each queue item is one complete serialized JSON message. +// subclass owns byte I/O and framing; base owns JSON and the JSON-RPC session (handshake, id correlation). +// + +struct server_mcp_transport { + std::string name; + int timeout_ms = 30000; + + server_pipe to_server; // serialized messages we send to the server + server_pipe from_server; // serialized messages read from the server + + virtual ~server_mcp_transport() = default; + + virtual bool start() = 0; + virtual void close() = 0; // blocking and idempotent + virtual bool is_alive() const = 0; // never blocks behind an in-flight send_rpc() + + // human-readable diagnostics for logging when the transport fails/dies + // (example: last RPC error, plus any transport-specific detail) + // may run on a different thread than send_rpc(), so last_error is read under rpc_mutex + virtual std::string diagnostics() { + std::lock_guard lock(rpc_mutex); + return last_error; + } + + std::vector list_tools(const std::function & should_stop); + + json call_tool(const std::string & tool_name, + const json & arguments, + const std::function & should_stop); + +protected: + // per-transport: send_rpc() holds it across the reply wait, so sharing it would stall every server behind one slow call. guards all members below. + std::mutex rpc_mutex; + uint64_t next_id = 1; // reset to 1 per (re)spawn + bool initialized = false; + std::string last_error; + std::vector tools; + + // both assume rpc_mutex is already held by the public caller + bool ensure_init(const std::function & should_stop); // initialize handshake, once + json send_rpc(const json & request, const std::function & should_stop); // returns the reply or an {"error": ...} +}; + +// +// server_mcp_stdio: child process, NDJSON JSON-RPC over stdio (stderr drained to the debug log) +// + +struct server_mcp_stdio : server_mcp_transport { + explicit server_mcp_stdio(const server_mcp_server_config & config); + ~server_mcp_stdio() override; + + bool start() override; + void close() override; + bool is_alive() const override; + std::string diagnostics() override; + +private: + server_mcp_server_config config; + + // defined in the .cpp so stays out of this header + struct process_handle; + std::unique_ptr proc; + + std::thread reader; // child stdout -> NDJSON de-framing -> from_server + std::thread writer; // to_server -> NDJSON framing -> child stdin + std::thread errlog; // child stderr -> debug log (must be drained or the child blocks) + + // cleared by close() or by the reader on stdout EOF; read without rpc_mutex + std::atomic running{false}; + + // bounded tail of the child's stderr, for diagnostics when it dies + std::mutex err_mu; + std::string err_tail; + + void reader_loop(); + void writer_loop(); + void errlog_loop(); + void join_pumps(); +}; + +// +// server_mcp +// declare before the HTTP context so it outlives every /tools handler. +// + +class server_mcp { +public: + server_mcp() = default; + ~server_mcp(); + + // parse the MCP config from params (file and/or inline JSON), + // then spawn each server once, list its tools, and shut it down + // throws on config parse errors; spawn failures are logged. + void start(const common_params & params); + + // true until start() has parsed at least one server from the config + bool empty() const { return configs.empty(); } + + std::vector list_tools() const; + + // lazily (re)spawns the transport. returns the MCP result or an {"error": ...}. should_stop is OR-ed with the manager's cancel flag. + json call_tool(const std::string & server_name, + const std::string & tool_name, + const json & arguments, + const std::function & should_stop = nullptr); + + // flip the cancel flag so in-flight calls return; blocking teardown is in the destructor. call before the HTTP server drains. + // note: multiple calls are idempotent + void shutdown(); + +private: + std::vector configs; + + mutable std::mutex mutex; // guards transports, dead_servers, registry + + // shared_ptr: call_tool() hands a transport to the caller and drops the lock for the blocking RPC, so a concurrent evict/respawn must not destroy it mid-call + std::map> transports; + std::map dead_servers; // spawn-failure cooldown + std::vector registry; + + std::atomic stopping{false}; + + const server_mcp_server_config * find_config(const std::string & name) const; + + // the only place that names a concrete transport + std::shared_ptr create_transport(const server_mcp_server_config & cfg); + + // nullptr during cooldown or shutdown + std::shared_ptr get_or_create(const std::string & name); +}; diff --git a/tools/server/server-models.cpp b/tools/server/server-models.cpp index 9eac58e9df..ba63788146 100644 --- a/tools/server/server-models.cpp +++ b/tools/server/server-models.cpp @@ -1944,53 +1944,6 @@ void server_models_routes::init_routes() { // server_http_proxy // -// simple implementation of a pipe -// used for streaming data between threads -template -struct pipe_t { - std::mutex mutex; - std::condition_variable cv; - std::queue queue; - std::atomic writer_closed{false}; - std::atomic reader_closed{false}; - void close_write() { - writer_closed.store(true, std::memory_order_relaxed); - cv.notify_all(); - } - void close_read() { - reader_closed.store(true, std::memory_order_relaxed); - cv.notify_all(); - } - bool read(T & output, const std::function & should_stop) { - std::unique_lock lk(mutex); - constexpr auto poll_interval = std::chrono::milliseconds(500); - while (true) { - if (!queue.empty()) { - output = std::move(queue.front()); - queue.pop(); - return true; - } - if (writer_closed.load()) { - return false; // clean EOF - } - if (should_stop()) { - close_read(); // signal broken pipe to writer - return false; // cancelled / reader no longer alive - } - cv.wait_for(lk, poll_interval); - } - } - bool write(T && data) { - std::lock_guard lk(mutex); - if (reader_closed.load()) { - return false; // broken pipe - } - queue.push(std::move(data)); - cv.notify_one(); - return true; - } -}; - static std::string to_lower_copy(const std::string & value) { std::string lowered(value.size(), '\0'); std::transform(value.begin(), value.end(), lowered.begin(), [](unsigned char c) { return std::tolower(c); }); @@ -2100,7 +2053,7 @@ server_http_proxy::server_http_proxy( ) { // shared between reader and writer threads auto cli = std::make_shared(host, port); - auto pipe = std::make_shared>(); + auto pipe = std::make_shared>(); if (scheme == "https") { #ifdef CPPHTTPLIB_OPENSSL_SUPPORT diff --git a/tools/server/server-tools.cpp b/tools/server/server-tools.cpp index a82a3d6025..2af44e49b9 100644 --- a/tools/server/server-tools.cpp +++ b/tools/server/server-tools.cpp @@ -10,10 +10,10 @@ #include #include #include -#include #include #include #include +#include namespace fs = std::filesystem; @@ -25,7 +25,7 @@ json server_tool::to_json() const { return { {"display_name", display_name}, {"tool", name}, - {"type", "builtin"}, + {"type", type()}, {"permissions", json{ {"write", permission_write} }}, @@ -1129,6 +1129,49 @@ struct server_tools_res : server_http_res { } }; +// +// server_mcp_tool: exposes one tool from a running MCP server as a server_tool. +// +struct server_mcp_tool : server_tool { + std::string server_name; + std::string tool_name; + server_mcp_tool_def def; + server_mcp & mcp_mgr; + + server_mcp_tool(server_mcp_tool_def d, server_mcp & mgr) + : server_name(d.server_name) + , tool_name(d.name) + , def(std::move(d)) + , mcp_mgr(mgr) + { + name = server_name + "_" + tool_name; + display_name = name; + permission_write = false; + support_stream = false; + } + + std::string type() const override { return "mcp"; } + + json get_definition() const override { + json schema = def.input_schema; + if (schema.is_null() || !schema.is_object()) { + schema = json::object(); + } + return { + {"type", "function"}, + {"function", { + {"name", name}, + {"description", def.description}, + {"parameters", schema}, + }}, + }; + } + + json invoke(json params, server_tool::stream *) const override { + return mcp_mgr.call_tool(server_name, tool_name, params); + } +}; + static server_tool & find_tool(std::vector> & tools, const std::string & name, bool require_stream) { for (auto & t : tools) { if (t->name == name) { @@ -1157,7 +1200,8 @@ static std::vector> build_tools() { return tools; } -void server_tools::setup(const std::vector & enabled_tools) { +void server_tools::setup(const std::vector & enabled_tools, + server_mcp & mcp_mgr) { if (!enabled_tools.empty()) { std::unordered_set enabled_set(enabled_tools.begin(), enabled_tools.end()); auto all_tools = build_tools(); @@ -1188,6 +1232,29 @@ void server_tools::setup(const std::vector & enabled_tools) { } } + // append MCP tools, skipping any that collide with a built-in or another MCP tool of the same "_" name + if (!mcp_mgr.empty()) { + std::unordered_set seen_names; + for (auto & t : tools) { + seen_names.insert(t->name); + } + size_t n_added = 0; + for (const auto & def : mcp_mgr.list_tools()) { + std::string mcp_name = def.server_name + "_" + def.name; + if (seen_names.count(mcp_name)) { + SRV_WRN("MCP tool \"%s\" from server \"%s\" collides with an existing tool, skipping\n", + mcp_name.c_str(), def.server_name.c_str()); + continue; + } + seen_names.insert(mcp_name); + tools.push_back(std::make_unique(def, mcp_mgr)); + n_added++; + } + if (n_added > 0) { + SRV_INF("Added %zu MCP tools\n", n_added); + } + } + handle_get = [this](const server_http_req &) -> server_http_res_ptr { auto res = std::make_unique(); try { diff --git a/tools/server/server-tools.h b/tools/server/server-tools.h index 6f6528f484..601399ee93 100644 --- a/tools/server/server-tools.h +++ b/tools/server/server-tools.h @@ -3,9 +3,11 @@ #include "server-common.h" #include "server-http.h" #include "server-queue.h" +#include "server-mcp.h" #include #include +#include struct server_tool { std::string name; @@ -15,6 +17,7 @@ struct server_tool { virtual ~server_tool() = default; virtual json get_definition() const = 0; + virtual std::string type() const { return "builtin"; } struct stream { server_response & qr; @@ -34,7 +37,8 @@ struct server_tools { server_response queue_res; std::atomic res_id{0}; - void setup(const std::vector & enabled_tools); + void setup(const std::vector & enabled_tools, + server_mcp & mcp_mgr); server_http_context::handler_t handle_get; server_http_context::handler_t handle_post; diff --git a/tools/server/server.cpp b/tools/server/server.cpp index 20effbb148..b6fef99e87 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -88,6 +88,11 @@ static server_http_context::handler_t ex_wrapper(server_http_context::handler_t int llama_server(int argc, char ** argv) { std::setlocale(LC_NUMERIC, "C"); +#ifndef _WIN32 + // Ignore SIGPIPE so the server does not crash if an MCP child exits while we are writing to its stdin + signal(SIGPIPE, SIG_IGN); +#endif + // own arguments required by this example common_params params; @@ -157,6 +162,9 @@ int llama_server(common_params & params, int argc, char ** argv) { params.model_alias.insert(model_name); } + // note: this is guaranteed to out-live ctx_http and tools + server_mcp mcp_mgr; + // struct that contains llama context and inference server_context ctx_server; @@ -326,17 +334,28 @@ int llama_server(common_params & params, int argc, char ** argv) { ctx_http.post("/cors-proxy", ex_wrapper(res_403)); } - // EXPERIMENTAL built-in tools - if (!params.server_tools.empty()) { + try { + mcp_mgr.start(params); + } catch (const std::exception & e) { + SRV_ERR("MCP starting failed: %s\n", e.what()); + return 1; + } + + if (!params.server_tools.empty() || !mcp_mgr.empty()) { try { - tools.setup(params.server_tools); + tools.setup(params.server_tools, mcp_mgr); } catch (const std::exception & e) { SRV_ERR("tools setup failed: %s\n", e.what()); return 1; } ctx_http.get ("/tools", ex_wrapper(tools.handle_get)); ctx_http.post("/tools", ex_wrapper(tools.handle_post)); - warn_names.push_back("built-in tools (experimental)"); + if (!params.server_tools.empty()) { + warn_names.push_back("built-in tools (experimental)"); + } + if (!mcp_mgr.empty()) { + warn_names.push_back("MCP servers (experimental)"); + } } else { ctx_http.get ("/tools", ex_wrapper(res_403)); ctx_http.post("/tools", ex_wrapper(res_403)); @@ -378,7 +397,7 @@ int llama_server(common_params & params, int argc, char ** argv) { if (is_router_server) { SRV_INF("%s", "starting server in router mode. models will be automatically loaded on-demand\n"); - clean_up = [&models_routes]() { + clean_up = [&models_routes, &mcp_mgr]() { SRV_INF("%s: cleaning up before exit...\n", __func__); // stop the session GC first, it finalizes live sessions and wakes pending readers server_stream_session_manager_stop(); @@ -386,6 +405,7 @@ int llama_server(common_params & params, int argc, char ** argv) { models_routes->stopping.store(true); // maybe redundant, but just to be safe models_routes->models.unload_all(); } + mcp_mgr.shutdown(); llama_backend_free(); }; @@ -401,17 +421,19 @@ int llama_server(common_params & params, int argc, char ** argv) { // important to disconnect any SSE clients models_routes->stopping.store(true); } + mcp_mgr.shutdown(); ctx_http.stop(); }; } else { // setup clean up function, to be called before exit - clean_up = [&ctx_http, &ctx_server]() { + clean_up = [&ctx_http, &ctx_server, &mcp_mgr]() { SRV_INF("%s: cleaning up before exit...\n", __func__); // stop the session GC first, it finalizes live sessions and wakes pending readers server_stream_session_manager_stop(); ctx_http.stop(); ctx_server.terminate(); + mcp_mgr.shutdown(); llama_backend_free(); }; @@ -444,6 +466,7 @@ int llama_server(common_params & params, int argc, char ** argv) { SRV_INF("%s", "model loaded\n"); shutdown_handler = [&](int) { + mcp_mgr.shutdown(); // this will unblock start_loop() ctx_server.terminate(); }; diff --git a/tools/server/tests/fixtures/mcp_burst_server.py b/tools/server/tests/fixtures/mcp_burst_server.py new file mode 100644 index 0000000000..22892d9a1d --- /dev/null +++ b/tools/server/tests/fixtures/mcp_burst_server.py @@ -0,0 +1,118 @@ +#!/usr/bin/env python3 +""" +Minimal MCP server that writes notification + response in a single write() with no flush. +This reproduces the buffering bug where read_message() can strand the response. +""" +import json +import sys +import os + +TOOLS = [ + { + "name": "echo", + "description": "Echo back the input message", + "inputSchema": { + "type": "object", + "properties": { + "message": {"type": "string"} + }, + "required": ["message"] + } + } +] + +def handle_initialize(params, req_id): + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "protocolVersion": "2024-11-05", + "capabilities": {"tools": {}}, + "serverInfo": {"name": "burst-test", "version": "1.0"} + } + } + +def handle_tools_list(params, req_id): + return { + "jsonrpc": "2.0", + "id": req_id, + "result": {"tools": TOOLS} + } + +def handle_tools_call(params, req_id): + tool_name = params.get("name") + arguments = params.get("arguments", {}) + + if tool_name == "echo": + message = arguments.get("message", "") + notif = { + "jsonrpc": "2.0", + "method": "notifications/progress", + "params": {"progress": 50, "total": 100} + } + response = { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "content": [{"type": "text", "text": f"echo: {message}"}] + } + } + # Single os.write() call: both lines land in one pipe packet atomically. + # This is the key difference from mcp_malformed_server.py which flushes between writes. + data = (json.dumps(notif) + "\n" + json.dumps(response) + "\n").encode("utf-8") + os.write(sys.stdout.fileno(), data) + return None # already written + else: + response = { + "jsonrpc": "2.0", + "id": req_id, + "error": {"code": -32602, "message": f"Unknown tool: {tool_name}"} + } + return response + +HANDLERS = { + "initialize": handle_initialize, + "tools/list": handle_tools_list, + "tools/call": handle_tools_call, +} + +def main(): + # Use line-buffered text mode for regular responses, but the burst write + # uses os.write() directly to guarantee a single kernel write(). + sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1) + sys.stderr = os.fdopen(sys.stderr.fileno(), "w", buffering=1) + + for line in sys.stdin: + line = line.strip() + if not line: + continue + try: + request = json.loads(line) + except json.JSONDecodeError: + continue + + method = request.get("method") + req_id = request.get("id") + params = request.get("params", {}) + + # JSON-RPC 2.0: a message without an id is a notification and must not receive a response + if req_id is None: + continue + + handler = HANDLERS.get(method) + if handler: + response = handler(params, req_id) + if response is not None: + sys.stdout.write(json.dumps(response) + "\n") + sys.stdout.flush() + else: + response = { + "jsonrpc": "2.0", + "id": req_id, + "error": {"code": -32601, "message": f"Method not found: {method}"} + } + sys.stdout.write(json.dumps(response) + "\n") + sys.stdout.flush() + +if __name__ == "__main__": + main() diff --git a/tools/server/tests/fixtures/mcp_crash_server.py b/tools/server/tests/fixtures/mcp_crash_server.py new file mode 100644 index 0000000000..8dffdc61c0 --- /dev/null +++ b/tools/server/tests/fixtures/mcp_crash_server.py @@ -0,0 +1,114 @@ +#!/usr/bin/env python3 +""" +MCP server that crashes after receiving a specific tool call. +""" +import json +import sys +import os + +def handle_initialize(params, req_id): + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "protocolVersion": "2024-11-05", + "capabilities": {"tools": {}}, + "serverInfo": {"name": "crash-test", "version": "1.0"} + } + } + +def handle_tools_list(params, req_id): + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "tools": [ + { + "name": "echo", + "description": "Echo back the input message", + "inputSchema": { + "type": "object", + "properties": { + "message": {"type": "string"} + } + } + }, + { + "name": "crash", + "description": "Crash the server", + "inputSchema": { + "type": "object", + "properties": {} + } + } + ] + } + } + +def handle_tools_call(params, req_id): + tool_name = params.get("name") + arguments = params.get("arguments", {}) + + if tool_name == "echo": + message = arguments.get("message", "") + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "content": [{"type": "text", "text": f"echo: {message}"}] + } + } + elif tool_name == "crash": + # Send a partial response then exit + sys.stdout.write(json.dumps({"jsonrpc": "2.0", "id": req_id, "result": {"content": [{"type": "text", "text": "crashing..."}]}}) + "\n") + sys.stdout.flush() + os._exit(1) + else: + return { + "jsonrpc": "2.0", + "id": req_id, + "error": {"code": -32602, "message": f"Unknown tool: {tool_name}"} + } + +HANDLERS = { + "initialize": handle_initialize, + "tools/list": handle_tools_list, + "tools/call": handle_tools_call, +} + +def main(): + sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1) + sys.stderr = os.fdopen(sys.stderr.fileno(), "w", buffering=1) + + for line in sys.stdin: + line = line.strip() + if not line: + continue + try: + request = json.loads(line) + except json.JSONDecodeError: + continue + + method = request.get("method") + req_id = request.get("id") + params = request.get("params", {}) + + # JSON-RPC 2.0: a message without an id is a notification and must not receive a response + if req_id is None: + continue + + handler = HANDLERS.get(method) + if handler: + response = handler(params, req_id) + else: + response = { + "jsonrpc": "2.0", + "id": req_id, + "error": {"code": -32601, "message": f"Method not found: {method}"} + } + + sys.stdout.write(json.dumps(response) + "\n") + sys.stdout.flush() + +if __name__ == "__main__": + main() diff --git a/tools/server/tests/fixtures/mcp_echo_server.py b/tools/server/tests/fixtures/mcp_echo_server.py new file mode 100755 index 0000000000..7acfb35888 --- /dev/null +++ b/tools/server/tests/fixtures/mcp_echo_server.py @@ -0,0 +1,164 @@ +#!/usr/bin/env python3 +""" +Minimal MCP server for testing. +Implements JSON-RPC 2.0 over stdio (line-delimited JSON). +""" +import json +import sys +import os + +# Ensure we use python3 from the current environment +if sys.platform == "win32": + # On Windows, we need to use the same python interpreter + pass + +TOOLS = [ + { + "name": "echo", + "description": "Echo back the input message", + "inputSchema": { + "type": "object", + "properties": { + "message": {"type": "string", "description": "Message to echo"} + }, + "required": ["message"] + } + }, + { + "name": "add", + "description": "Add two numbers", + "inputSchema": { + "type": "object", + "properties": { + "a": {"type": "number"}, + "b": {"type": "number"} + }, + "required": ["a", "b"] + } + }, + { + "name": "fail_once", + "description": "Fails on first call, succeeds on subsequent calls", + "inputSchema": { + "type": "object", + "properties": {} + } + } +] + +_state = {"fail_once_called": False} + +def handle_initialize(params, req_id): + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "protocolVersion": "2024-11-05", + "capabilities": {"tools": {}}, + "serverInfo": {"name": "echo-test", "version": "1.0"} + } + } + +def handle_tools_list(params, req_id): + return { + "jsonrpc": "2.0", + "id": req_id, + "result": {"tools": TOOLS} + } + +def handle_tools_call(params, req_id): + tool_name = params.get("name") + arguments = params.get("arguments", {}) + + if tool_name == "echo": + message = arguments.get("message", "") + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "content": [{"type": "text", "text": f"echo: {message}"}] + } + } + elif tool_name == "add": + a = arguments.get("a", 0) + b = arguments.get("b", 0) + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "content": [{"type": "text", "text": str(a + b)}] + } + } + elif tool_name == "fail_once": + if not _state["fail_once_called"]: + _state["fail_once_called"] = True + return { + "jsonrpc": "2.0", + "id": req_id, + "error": {"code": -32000, "message": "transient error"} + } + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "content": [{"type": "text", "text": "ok"}] + } + } + else: + return { + "jsonrpc": "2.0", + "id": req_id, + "error": {"code": -32602, "message": f"Unknown tool: {tool_name}"} + } + +def handle_ping(params, req_id): + return { + "jsonrpc": "2.0", + "id": req_id, + "result": {} + } + +HANDLERS = { + "initialize": handle_initialize, + "tools/list": handle_tools_list, + "tools/call": handle_tools_call, + "ping": handle_ping, +} + +def main(): + # Use unbuffered output + sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1) + sys.stderr = os.fdopen(sys.stderr.fileno(), "w", buffering=1) + + for line in sys.stdin: + line = line.strip() + if not line: + continue + try: + request = json.loads(line) + except json.JSONDecodeError: + continue + + method = request.get("method") + req_id = request.get("id") + params = request.get("params", {}) + + # JSON-RPC 2.0: a message without an id is a notification and must not receive a response + if req_id is None: + continue + + handler = HANDLERS.get(method) + if handler: + response = handler(params, req_id) + else: + response = { + "jsonrpc": "2.0", + "id": req_id, + "error": {"code": -32601, "message": f"Method not found: {method}"} + } + + sys.stdout.write(json.dumps(response) + "\n") + sys.stdout.flush() + +if __name__ == "__main__": + main() diff --git a/tools/server/tests/fixtures/mcp_grandchild_server.py b/tools/server/tests/fixtures/mcp_grandchild_server.py new file mode 100644 index 0000000000..2604a77ee8 --- /dev/null +++ b/tools/server/tests/fixtures/mcp_grandchild_server.py @@ -0,0 +1,100 @@ +#!/usr/bin/env python3 +""" +MCP server (NDJSON JSON-RPC over stdio) that spawns a long-lived grandchild which inherits +this process's stdin/stdout/stderr and keeps them open. + +This reproduces the reader-teardown deadlock: killing the direct MCP child (SIGKILL, which is +all subprocess_terminate() does) does NOT close the stdout/stderr pipe write ends, because the +grandchild still holds them. A server that reads those pipes with a blocking read would then +wait forever for an EOF that never arrives, hanging teardown (both warmup shutdown at startup +and process shutdown). The polled, running-aware reader must exit regardless. +""" +import json +import os +import subprocess +import sys + +# Spawn a grandchild that inherits our std handles (fds 0/1/2 = the MCP pipes) and lives well +# past any teardown in the tests. We do NOT redirect its stdio, so it keeps the pipe write ends +# open even after this process is killed. +subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"]) + +TOOLS = [ + { + "name": "echo", + "description": "Echo back the input message", + "inputSchema": { + "type": "object", + "properties": {"message": {"type": "string", "description": "Message to echo"}}, + "required": ["message"], + }, + } +] + + +def handle_initialize(params, req_id): + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "protocolVersion": "2024-11-05", + "capabilities": {"tools": {}}, + "serverInfo": {"name": "grandchild-test", "version": "1.0"}, + }, + } + + +def handle_tools_list(params, req_id): + return {"jsonrpc": "2.0", "id": req_id, "result": {"tools": TOOLS}} + + +def handle_tools_call(params, req_id): + if params.get("name") == "echo": + message = params.get("arguments", {}).get("message", "") + return { + "jsonrpc": "2.0", + "id": req_id, + "result": {"content": [{"type": "text", "text": f"echo: {message}"}]}, + } + return {"jsonrpc": "2.0", "id": req_id, "error": {"code": -32602, "message": "Unknown tool"}} + + +HANDLERS = { + "initialize": handle_initialize, + "tools/list": handle_tools_list, + "tools/call": handle_tools_call, +} + + +def main(): + sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1) + sys.stderr = os.fdopen(sys.stderr.fileno(), "w", buffering=1) + + for line in sys.stdin: + line = line.strip() + if not line: + continue + try: + request = json.loads(line) + except json.JSONDecodeError: + continue + + method = request.get("method") + req_id = request.get("id") + params = request.get("params", {}) + + if req_id is None: + continue # notification, no response + + handler = HANDLERS.get(method) + if handler: + response = handler(params, req_id) + else: + response = {"jsonrpc": "2.0", "id": req_id, "error": {"code": -32601, "message": f"Method not found: {method}"}} + + sys.stdout.write(json.dumps(response) + "\n") + sys.stdout.flush() + + +if __name__ == "__main__": + main() diff --git a/tools/server/tests/fixtures/mcp_malformed_server.py b/tools/server/tests/fixtures/mcp_malformed_server.py new file mode 100644 index 0000000000..743333c5fd --- /dev/null +++ b/tools/server/tests/fixtures/mcp_malformed_server.py @@ -0,0 +1,113 @@ +#!/usr/bin/env python3 +""" +MCP server that sends malformed responses and notifications during requests. +""" +import json +import sys +import os + +def handle_initialize(params, req_id): + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "protocolVersion": "2024-11-05", + "capabilities": {"tools": {}}, + "serverInfo": {"name": "malformed-test", "version": "1.0"} + } + } + +def handle_tools_list(params, req_id): + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "tools": [ + { + "name": "echo", + "description": "Echo back the input message", + "inputSchema": { + "type": "object", + "properties": { + "message": {"type": "string"} + } + } + } + ] + } + } + +def handle_tools_call(params, req_id): + tool_name = params.get("name") + arguments = params.get("arguments", {}) + + if tool_name == "echo": + message = arguments.get("message", "") + # Send a notification first (no id field) + notif = { + "jsonrpc": "2.0", + "method": "notifications/progress", + "params": {"progress": 50, "total": 100} + } + sys.stdout.write(json.dumps(notif) + "\n") + sys.stdout.flush() + # Then send the actual response + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "content": [{"type": "text", "text": f"echo: {message}"}] + } + } + else: + return { + "jsonrpc": "2.0", + "id": req_id, + "error": {"code": -32602, "message": f"Unknown tool: {tool_name}"} + } + +HANDLERS = { + "initialize": handle_initialize, + "tools/list": handle_tools_list, + "tools/call": handle_tools_call, +} + +def main(): + sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1) + sys.stderr = os.fdopen(sys.stderr.fileno(), "w", buffering=1) + + for line in sys.stdin: + line = line.strip() + if not line: + continue + try: + request = json.loads(line) + except json.JSONDecodeError: + # Send malformed JSON response + sys.stdout.write("THIS IS NOT JSON\n") + sys.stdout.flush() + continue + + method = request.get("method") + req_id = request.get("id") + params = request.get("params", {}) + + # JSON-RPC 2.0: a message without an id is a notification and must not receive a response + if req_id is None: + continue + + handler = HANDLERS.get(method) + if handler: + response = handler(params, req_id) + else: + response = { + "jsonrpc": "2.0", + "id": req_id, + "error": {"code": -32601, "message": f"Method not found: {method}"} + } + + sys.stdout.write(json.dumps(response) + "\n") + sys.stdout.flush() + +if __name__ == "__main__": + main() diff --git a/tools/server/tests/fixtures/mcp_slow_server.py b/tools/server/tests/fixtures/mcp_slow_server.py new file mode 100644 index 0000000000..7f8e67835a --- /dev/null +++ b/tools/server/tests/fixtures/mcp_slow_server.py @@ -0,0 +1,132 @@ +#!/usr/bin/env python3 +""" +MCP server that sleeps before responding, for timeout testing. +""" +import json +import sys +import os +import time +import argparse + +TOOLS = [ + { + "name": "sleep", + "description": "Sleep for a given number of seconds", + "inputSchema": { + "type": "object", + "properties": { + "seconds": {"type": "number", "description": "Seconds to sleep"} + }, + "required": ["seconds"] + } + } +] + +def handle_initialize(params, req_id): + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "protocolVersion": "2024-11-05", + "capabilities": {"tools": {}}, + "serverInfo": {"name": "slow-test", "version": "1.0"} + } + } + +def handle_tools_list(params, req_id): + return { + "jsonrpc": "2.0", + "id": req_id, + "result": {"tools": TOOLS} + } + +def handle_tools_call(params, req_id): + tool_name = params.get("name") + arguments = params.get("arguments", {}) + + if tool_name == "sleep": + seconds = arguments.get("seconds", 1) + time.sleep(seconds) + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "content": [{"type": "text", "text": f"slept {seconds}s"}] + } + } + else: + return { + "jsonrpc": "2.0", + "id": req_id, + "error": {"code": -32602, "message": f"Unknown tool: {tool_name}"} + } + +HANDLERS = { + "initialize": handle_initialize, + "tools/list": handle_tools_list, + "tools/call": handle_tools_call, +} + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--delay", type=float, default=5.0, help="Delay in seconds for sleep tool") + args = parser.parse_args() + + # Override the sleep duration + global handle_tools_call + def handle_tools_call(params, req_id): + tool_name = params.get("name") + arguments = params.get("arguments", {}) + + if tool_name == "sleep": + seconds = arguments.get("seconds", args.delay) + time.sleep(seconds) + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "content": [{"type": "text", "text": f"slept {seconds}s"}] + } + } + else: + return { + "jsonrpc": "2.0", + "id": req_id, + "error": {"code": -32602, "message": f"Unknown tool: {tool_name}"} + } + + sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1) + sys.stderr = os.fdopen(sys.stderr.fileno(), "w", buffering=1) + + for line in sys.stdin: + line = line.strip() + if not line: + continue + try: + request = json.loads(line) + except json.JSONDecodeError: + continue + + method = request.get("method") + req_id = request.get("id") + params = request.get("params", {}) + + # JSON-RPC 2.0: a message without an id is a notification and must not receive a response + if req_id is None: + continue + + handler = HANDLERS.get(method) + if handler: + response = handler(params, req_id) + else: + response = { + "jsonrpc": "2.0", + "id": req_id, + "error": {"code": -32601, "message": f"Method not found: {method}"} + } + + sys.stdout.write(json.dumps(response) + "\n") + sys.stdout.flush() + +if __name__ == "__main__": + main() diff --git a/tools/server/tests/unit/test_mcp_servers.py b/tools/server/tests/unit/test_mcp_servers.py new file mode 100644 index 0000000000..9ad2241bd0 --- /dev/null +++ b/tools/server/tests/unit/test_mcp_servers.py @@ -0,0 +1,718 @@ +#!/usr/bin/env python3 +""" +Tests for MCP server integration via the /tools endpoint. + +Invariants verified: +1. MCP tools appear in /tools listing when configured +2. MCP tools use _ naming +3. MCP tools can be invoked and return correct results +4. Misconfigured MCP servers do not crash the server +5. Multiple MCP servers can be configured simultaneously +6. Warmup populates the tool list at startup +""" +import json +import os +import sys +import tempfile +import time + +import pytest + +from utils import * + +# Path to the test MCP server fixture +FIXTURES_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "fixtures") +MCP_ECHO_SERVER = os.path.join(FIXTURES_DIR, "mcp_echo_server.py") + +server: ServerProcess + + +def _mcp_config_json(servers: dict) -> str: + """Create a JSON config string for --mcp-servers-json.""" + return json.dumps({"mcpServers": servers}) + + +def _start_server_with_mcp(mcp_json: str, **kwargs) -> ServerProcess: + """Helper to start a router server with MCP config.""" + srv = ServerPreset.router() + srv.server_tools = "all" + srv.no_ui = True + srv.server_port = 8085 # avoid conflict with load_all() which uses 8080 + srv.mcp_servers_json = mcp_json + for k, v in kwargs.items(): + setattr(srv, k, v) + srv.start() + return srv + + +def test_mcp_tools_listed_in_tools_endpoint(): + """MCP tools should appear in GET /tools with server:tool naming.""" + global server + mcp_json = _mcp_config_json({ + "echo": { + "command": sys.executable, + "args": [MCP_ECHO_SERVER], + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + res = server.make_request("GET", "/tools") + assert res.status_code == 200, res.body + + tools = res.body + assert isinstance(tools, list), f"Expected list, got {type(tools)}" + + # Find MCP tools - name is in "tool" field or definition.function.name + def get_tool_name(t): + return t.get("tool", "") or t.get("definition", {}).get("function", {}).get("name", "") + + mcp_tools = [t for t in tools if get_tool_name(t).startswith("echo_")] + assert len(mcp_tools) >= 2, f"Expected at least 2 echo_ tools, got {len(mcp_tools)}: {mcp_tools}" + + tool_names = {get_tool_name(t) for t in mcp_tools} + assert "echo_echo" in tool_names + assert "echo_add" in tool_names + + # Verify tool structure + echo_tool = next(t for t in mcp_tools if get_tool_name(t) == "echo_echo") + assert "description" in echo_tool or "definition" in echo_tool + finally: + server.stop() + + +def test_mcp_tool_invocation(): + """MCP tools should be callable via POST /tools and return correct results.""" + global server + mcp_json = _mcp_config_json({ + "echo": { + "command": sys.executable, + "args": [MCP_ECHO_SERVER], + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + # Call echo_echo + res = server.make_request("POST", "/tools", data={ + "tool": "echo_echo", + "params": {"message": "hello world"} + }) + assert res.status_code == 200, res.body + body = res.body + assert "error" not in body, body + # The result format depends on the tool implementation + # For MCP tools, it should contain the tool result + assert "plain_text_response" in body or "result" in body or "content" in body, body + + # Call echo_add + res = server.make_request("POST", "/tools", data={ + "tool": "echo_add", + "params": {"a": 3, "b": 5} + }) + assert res.status_code == 200, res.body + body = res.body + assert "error" not in body, body + finally: + server.stop() + + +def test_mcp_bad_command_does_not_crash(): + """A misconfigured MCP server should not crash the llama-server.""" + global server + mcp_json = _mcp_config_json({ + "nonexistent": { + "command": "this_executable_does_not_exist_12345", + "args": [], + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + # Server should still be healthy + res = server.make_request("GET", "/health") + assert res.status_code == 200, res.body + + # Builtin tools should still work + res = server.make_request("GET", "/tools") + assert res.status_code == 200, res.body + tools = res.body + # Should have builtin tools but no MCP tools from the bad server + mcp_tools = [t for t in tools if t.get("name", "").startswith("nonexistent_")] + assert len(mcp_tools) == 0, f"Expected no nonexistent_ tools, got {mcp_tools}" + finally: + server.stop() + + +def test_mcp_multiple_servers(): + """Multiple MCP servers can be configured simultaneously.""" + global server + mcp_json = _mcp_config_json({ + "echo": { + "command": sys.executable, + "args": [MCP_ECHO_SERVER], + }, + "echo2": { + "command": sys.executable, + "args": [MCP_ECHO_SERVER], + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + res = server.make_request("GET", "/tools") + assert res.status_code == 200, res.body + + tools = res.body + + def get_tool_name(t): + return t.get("tool", "") or t.get("definition", {}).get("function", {}).get("name", "") + + echo_tools = [t for t in tools if get_tool_name(t).startswith("echo_")] + echo2_tools = [t for t in tools if get_tool_name(t).startswith("echo2_")] + + assert len(echo_tools) >= 2, f"Expected echo_ tools, got {echo_tools}" + assert len(echo2_tools) >= 2, f"Expected echo2_ tools, got {echo2_tools}" + finally: + server.stop() + + +def test_mcp_tools_not_listed_when_not_configured(): + """Without MCP config, no MCP tools should appear.""" + global server + server = ServerPreset.router() + server.server_tools = "all" + server.no_ui = True + server.server_port = 8085 + server.start() + + try: + res = server.make_request("GET", "/tools") + assert res.status_code == 200, res.body + + tools = res.body + + def get_tool_name(t): + return t.get("tool", "") or t.get("definition", {}).get("function", {}).get("name", "") + + # Should only have builtin tools, no server: prefixed tools + mcp_tools = [t for t in tools if ":" in get_tool_name(t)] + assert len(mcp_tools) == 0, f"Expected no MCP tools, got {mcp_tools}" + finally: + server.stop() + + +def test_mcp_fail_once_tool_eventual_success(): + """Test that a tool that fails once eventually succeeds (tests instance respawn).""" + global server + mcp_json = _mcp_config_json({ + "echo": { + "command": sys.executable, + "args": [MCP_ECHO_SERVER], + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + # First call should succeed (warmup already spawned and shut down the instance, + # but the first actual tool call will spawn a fresh instance) + res = server.make_request("POST", "/tools", data={ + "tool": "echo_fail_once", + "params": {} + }) + # It might fail on first call if the warmup instance was shut down + # and a new instance is spawned. The fail_once state is per-process, + # so a fresh process will fail once then succeed. + # Actually, warmup spawns, lists, then shuts down. So the first tool call + # spawns a new process which will fail once. + assert res.status_code in (200, 500), res.body + finally: + server.stop() + + +def test_mcp_tools_via_json_config_file(): + """Test that --mcp-servers-config (file) works as well as --mcp-servers-json.""" + global server + config = { + "mcpServers": { + "echo": { + "command": sys.executable, + "args": [MCP_ECHO_SERVER], + } + } + } + + with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: + json.dump(config, f) + config_path = f.name + + try: + server = ServerPreset.router() + server.server_tools = "all" + server.no_ui = True + server.server_port = 8085 + server.mcp_servers_config = config_path + server.start() + + res = server.make_request("GET", "/tools") + assert res.status_code == 200, res.body + + tools = res.body + + def get_tool_name(t): + return t.get("tool", "") or t.get("definition", {}).get("function", {}).get("name", "") + + mcp_tools = [t for t in tools if get_tool_name(t).startswith("echo_")] + assert len(mcp_tools) >= 2, f"Expected echo_ tools, got {mcp_tools}" + finally: + os.unlink(config_path) + server.stop() + + +def test_mcp_tools_slot_independent(): + """MCP tools should work without any slot concept; /tools is slot-independent.""" + global server + mcp_json = _mcp_config_json({ + "echo": { + "command": sys.executable, + "args": [MCP_ECHO_SERVER], + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + # Call /tools without any slot binding - should succeed + res = server.make_request("POST", "/tools", data={ + "tool": "echo_echo", + "params": {"message": "hello"} + }) + assert res.status_code == 200, res.body + body = res.body + assert "error" not in body, body + finally: + server.stop() + + +def test_mcp_concurrent_tool_calls(): + """Concurrent POST /tools to same MCP server should all succeed.""" + global server + mcp_json = _mcp_config_json({ + "echo": { + "command": sys.executable, + "args": [MCP_ECHO_SERVER], + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + def call_tool(): + return server.make_request("POST", "/tools", data={ + "tool": "echo_echo", + "params": {"message": "hi"} + }) + + with ThreadPoolExecutor(max_workers=10) as executor: + futures = [executor.submit(call_tool) for _ in range(10)] + results = [f.result() for f in futures] + + for res in results: + assert res.status_code == 200, res.body + assert "error" not in res.body, res.body + finally: + server.stop() + + +def test_mcp_tool_timeout(): + """Tool call should timeout if MCP server is too slow.""" + global server + MCP_SLOW_SERVER = os.path.join(FIXTURES_DIR, "mcp_slow_server.py") + mcp_json = _mcp_config_json({ + "slow": { + "command": sys.executable, + "args": [MCP_SLOW_SERVER, "--delay", "5"], + "timeout_ms": 500 + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + res = server.make_request("POST", "/tools", data={ + "tool": "slow_sleep", + "params": {"seconds": 5} + }) + assert res.status_code == 200, res.body + body = res.body + assert "error" in body, body + finally: + server.stop() + + +def test_mcp_warmup_partial_failure(): + """Good server's tools should appear even if bad server fails warmup.""" + global server + mcp_json = _mcp_config_json({ + "good": { + "command": sys.executable, + "args": [MCP_ECHO_SERVER], + }, + "bad": { + "command": "nonexistent", + "args": [] + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + res = server.make_request("GET", "/tools") + assert res.status_code == 200, res.body + tools = res.body + + def get_tool_name(t): + return t.get("tool", "") or t.get("definition", {}).get("function", {}).get("name", "") + + # good server tools should be present + assert any("good_" in get_tool_name(t) for t in tools), f"Expected good: tools in {tools}" + finally: + server.stop() + + +def test_mcp_notification_during_request(): + """Notification during request should not be returned as response.""" + global server + MCP_MALFORMED_SERVER = os.path.join(FIXTURES_DIR, "mcp_malformed_server.py") + mcp_json = _mcp_config_json({ + "notifying": { + "command": sys.executable, + "args": [MCP_MALFORMED_SERVER], + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + res = server.make_request("POST", "/tools", data={ + "tool": "notifying_echo", + "params": {"message": "hi"} + }) + assert res.status_code == 200, res.body + body = res.body + assert "error" not in body, body + finally: + server.stop() + + +def test_mcp_instance_respawn_after_crash(): + """Tool call after process crash should respawn and succeed.""" + global server + MCP_CRASH_SERVER = os.path.join(FIXTURES_DIR, "mcp_crash_server.py") + mcp_json = _mcp_config_json({ + "crash": { + "command": sys.executable, + "args": [MCP_CRASH_SERVER], + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + # First call succeeds + res1 = server.make_request("POST", "/tools", data={ + "tool": "crash_echo", + "params": {"message": "hi"} + }) + assert res1.status_code == 200, res1.body + assert "error" not in res1.body, res1.body + + # Second call should also succeed (respawned instance) + res2 = server.make_request("POST", "/tools", data={ + "tool": "crash_echo", + "params": {"message": "hi2"} + }) + assert res2.status_code == 200, res2.body + assert "error" not in res2.body, res2.body + finally: + server.stop() + + + + +def test_mcp_fail_once_eventual_success_verified(): + """Verify that fail_once tool eventually succeeds after respawn.""" + global server + mcp_json = _mcp_config_json({ + "echo": { + "command": sys.executable, + "args": [MCP_ECHO_SERVER], + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + # First call may fail (fresh process) + res1 = server.make_request("POST", "/tools", data={ + "tool": "echo_fail_once", + "params": {} + }) + # Second call should succeed + res2 = server.make_request("POST", "/tools", data={ + "tool": "echo_fail_once", + "params": {} + }) + assert res2.status_code == 200, res2.body + assert "error" not in res2.body, res2.body + finally: + server.stop() + + +def test_mcp_config_file_errors(): + """Invalid JSON config and missing file should cause server to fail to start.""" + # Invalid JSON - server should fail to start + server = ServerPreset.router() + server.server_tools = "all" + server.no_ui = True + server.server_port = 8085 + server.mcp_servers_json = "not valid json" + try: + server.start() + assert False, "Server should not have started with invalid MCP JSON config" + except RuntimeError: + pass # Expected: server process dies due to bad config + + # Missing file - server should fail to start + server = ServerPreset.router() + server.server_tools = "all" + server.no_ui = True + server.server_port = 8085 + server.mcp_servers_config = "/nonexistent/path.json" + try: + server.start() + assert False, "Server should not have started with missing config file" + except RuntimeError: + pass # Expected: server process dies due to missing config + + +def test_mcp_empty_tool_list(): + """MCP server reporting zero tools should result in empty tool list.""" + global server + # Create a minimal server that returns empty tools list + empty_server = os.path.join(FIXTURES_DIR, "_empty_mcp_server.py") + with open(empty_server, "w") as f: + f.write('''#!/usr/bin/env python3 +import json, sys, os +def main(): + sys.stdout = os.fdopen(sys.stdout.fileno(), "w", buffering=1) + for line in sys.stdin: + line = line.strip() + if not line: continue + try: request = json.loads(line) + except: continue + method = request.get("method") + req_id = request.get("id") + if method == "initialize": + resp = {"jsonrpc": "2.0", "id": req_id, "result": {"protocolVersion": "2024-11-05", "capabilities": {"tools": {}}, "serverInfo": {"name": "empty", "version": "1.0"}}} + elif method == "tools/list": + resp = {"jsonrpc": "2.0", "id": req_id, "result": {"tools": []}} + else: + resp = {"jsonrpc": "2.0", "id": req_id, "error": {"code": -32601, "message": "Method not found"}} + sys.stdout.write(json.dumps(resp) + "\\n") + sys.stdout.flush() +if __name__ == "__main__": + main() +''') + try: + mcp_json = _mcp_config_json({ + "empty": { + "command": sys.executable, + "args": [empty_server], + } + }) + server = _start_server_with_mcp(mcp_json) + res = server.make_request("GET", "/tools") + assert res.status_code == 200, res.body + tools = res.body + def get_tool_name(t): + return t.get("tool", "") or t.get("definition", {}).get("function", {}).get("name", "") + mcp_tools = [t for t in tools if get_tool_name(t).startswith("empty:")] + assert len(mcp_tools) == 0, f"Expected no empty: tools, got {mcp_tools}" + finally: + os.unlink(empty_server) + server.stop() + + +def test_mcp_rapid_succession_calls(): + """Many rapid calls should increment next_id correctly and correlate responses.""" + global server + mcp_json = _mcp_config_json({ + "echo": { + "command": sys.executable, + "args": [MCP_ECHO_SERVER], + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + for i in range(20): + res = server.make_request("POST", "/tools", data={ + "tool": "echo_echo", + "params": {"message": f"msg{i}"} + }) + assert res.status_code == 200, res.body + assert "error" not in res.body, res.body + finally: + server.stop() + + +def test_mcp_notification_burst(): + """Notification + response in a single write() with no flush should not strand the response.""" + global server + MCP_BURST_SERVER = os.path.join(FIXTURES_DIR, "mcp_burst_server.py") + mcp_json = _mcp_config_json({ + "burst": { + "command": sys.executable, + "args": [MCP_BURST_SERVER], + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + res = server.make_request("POST", "/tools", data={ + "tool": "burst_echo", + "params": {"message": "burst test"} + }) + assert res.status_code == 200, res.body + body = res.body + assert "error" not in body, body + finally: + server.stop() + + +def test_mcp_tool_definition_shape_via_chat_completions(): + """MCP tool definitions returned by GET /tools should have the correct shape for chat/completions.""" + global server + mcp_json = _mcp_config_json({ + "echo": { + "command": sys.executable, + "args": [MCP_ECHO_SERVER], + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + # Get MCP tool definitions + res = server.make_request("GET", "/tools") + assert res.status_code == 200, res.body + tools = res.body + + def get_tool_name(t): + return t.get("tool", "") or t.get("definition", {}).get("function", {}).get("name", "") + + echo_tools = [t for t in tools if get_tool_name(t).startswith("echo_")] + assert len(echo_tools) >= 2, f"Expected echo_ tools, got {echo_tools}" + + echo_tool = next(t for t in echo_tools if get_tool_name(t) == "echo_echo") + definition = echo_tool.get("definition", echo_tool) + + # Verify the definition has the standard function-calling shape + assert definition.get("type") == "function", f"Expected type=function, got {definition.get('type')}" + func = definition.get("function", {}) + assert "name" in func, "Missing function.name" + assert "description" in func, "Missing function.description" + assert "parameters" in func, f"Missing function.parameters, got keys: {list(func.keys())}" + params = func["parameters"] + assert params.get("type") == "object", f"Expected parameters.type=object, got {params.get('type')}" + assert "properties" in params, "Missing parameters.properties" + finally: + server.stop() + + +def test_mcp_slow_tool_call_slot_release(): + """A slow tool call should not stall server shutdown for the full I/O timeout.""" + global server + MCP_SLOW_SERVER = os.path.join(FIXTURES_DIR, "mcp_slow_server.py") + mcp_json = _mcp_config_json({ + "slow": { + "command": sys.executable, + "args": [MCP_SLOW_SERVER, "--delay", "10"], + "timeout_ms": 30000 + } + }) + server = _start_server_with_mcp(mcp_json) + + try: + # Start a slow tool call in a background thread + def slow_call(): + return server.make_request("POST", "/tools", data={ + "tool": "slow_sleep", + "params": {"seconds": 10} + }) + + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(slow_call) + + # Wait a moment for the call to start + time.sleep(2) + + # Stop the server while the tool call is in progress. + # With global MCP instances, close_all() is called explicitly at shutdown + # (not from slot release), so shutdown should complete promptly. + start_time = time.time() + server.stop() + elapsed = time.time() - start_time + + # The server should stop quickly, not wait for the full 30s I/O timeout. + # With the terminating flag, send_rpc() bails out within one select() + # slice (~50ms). This threshold MUST stay below the 5s force-kill + # fallback in ServerProcess.stop(): without the flag, shutdown stalls + # on the instance mutex and only completes when stop() sends SIGKILL + # at ~5s -- which any threshold above 5 would still accept. + assert elapsed < 3, f"Server stop took {elapsed:.1f}s, expected < 3s" + + # Wait for the future to complete (it will get an error response or timeout) + try: + res = future.result(timeout=5) + # If we got a response, it should be an error since the server stopped + if hasattr(res, 'status_code'): + assert res.status_code in (200, 500, 502, 503, 504), f"Unexpected status: {res.status_code}" + except Exception: + # Thread may have raised due to connection error - that's acceptable + pass + finally: + server.stop() + + +def test_mcp_grandchild_holding_pipes_does_not_deadlock(): + """An MCP server that leaves a grandchild inheriting its stdout/stderr must not deadlock + teardown. + + subprocess_terminate() only SIGKILLs the direct MCP child, so the inherited pipe write ends + stay open and a blocking read on them would never see EOF. That hung both warmup shutdown + (the server would never reach "ready") and process shutdown. The polled, running-aware reader + must exit regardless, so the server both starts and stops promptly here. + """ + global server + MCP_GRANDCHILD_SERVER = os.path.join(FIXTURES_DIR, "mcp_grandchild_server.py") + mcp_json = _mcp_config_json({ + "gc": { + "command": sys.executable, + "args": [MCP_GRANDCHILD_SERVER], + } + }) + + # If warmup teardown deadlocked, the server would never become ready and start() would time out. + server = _start_server_with_mcp(mcp_json) + + try: + # invoking the tool spawns a live transport whose reader thread holds the inherited pipe + res = server.make_request("POST", "/tools", data={ + "tool": "gc_echo", + "params": {"message": "hello"} + }) + assert res.status_code == 200, res.body + assert "error" not in res.body, res.body + + # shutdown must be prompt: a deadlocked reader-join would stall until the 5s SIGKILL + # fallback in ServerProcess.stop(), so the threshold has to stay below that + start = time.time() + server.stop() + elapsed = time.time() - start + assert elapsed < 3, f"server shutdown took {elapsed:.1f}s (expected < 3s) — teardown likely deadlocked" + finally: + server.stop() diff --git a/tools/server/tests/utils.py b/tools/server/tests/utils.py index 5d5c873ac4..ae56bc70a1 100644 --- a/tools/server/tests/utils.py +++ b/tools/server/tests/utils.py @@ -115,6 +115,8 @@ class ServerProcess: backend_sampling: bool = False gcp_compat: bool = False server_tools: str | None = None + mcp_servers_config: str | None = None + mcp_servers_json: str | None = None cors_origins: str | None = None # session variables @@ -265,6 +267,10 @@ class ServerProcess: server_args.append("--ui-mcp-proxy") if self.server_tools: server_args.extend(["--tools", self.server_tools]) + if self.mcp_servers_config: + server_args.extend(["--mcp-servers-config", self.mcp_servers_config]) + if self.mcp_servers_json: + server_args.extend(["--mcp-servers-json", self.mcp_servers_json]) if self.backend_sampling: server_args.append("--backend_sampling") if self.gcp_compat: