diff --git a/common/CMakeLists.txt b/common/CMakeLists.txt index dd8d06e8f1..799d227519 100644 --- a/common/CMakeLists.txt +++ b/common/CMakeLists.txt @@ -102,6 +102,8 @@ add_library(${TARGET} speculative.h subproc.cpp subproc.h + trie.cpp + trie.h unicode.cpp unicode.h jinja/lexer.cpp 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/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..2792521836 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; @@ -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/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/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; 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/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); 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; 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/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/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; } 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.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-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 dfa9c467ee..923b3533e9 100644 --- a/tools/server/server-models.cpp +++ b/tools/server/server-models.cpp @@ -1910,53 +1910,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); }); @@ -2066,7 +2019,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-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(), 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/server-tools.cpp b/tools/server/server-tools.cpp index 8589e3d7c1..90b7e2a9f0 100644 --- a/tools/server/server-tools.cpp +++ b/tools/server/server-tools.cpp @@ -7,12 +7,13 @@ #include #include #include +#include #include #include -#include #include #include #include +#include namespace fs = std::filesystem; @@ -24,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} }}, @@ -1036,16 +1037,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)}}; } }; @@ -1090,6 +1117,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) { @@ -1118,7 +1188,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()) { if (!common_subproc::is_supported()) { throw std::runtime_error("subprocess is not enabled on this build"); @@ -1153,6 +1224,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_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 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: 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;