diff --git a/common/chat-auto-parser-generator.cpp b/common/chat-auto-parser-generator.cpp index 36aab7ecb..ddf81fc4d 100644 --- a/common/chat-auto-parser-generator.cpp +++ b/common/chat-auto-parser-generator.cpp @@ -147,7 +147,8 @@ common_peg_arena autoparser::build_parser(const generation_params & inputs, cons } else { parser = content.build_parser(ctx); } - return pure_content ? p.prefix(generation_prompt, reasoning.start) + parser : p.prefix(generation_prompt, reasoning.start) << parser; + const std::string reasoning_start = trim_whitespace(reasoning.start); + return pure_content ? p.prefix(generation_prompt, reasoning_start) + parser : p.prefix(generation_prompt, reasoning_start) << parser; }); } diff --git a/common/chat-diff-analyzer.cpp b/common/chat-diff-analyzer.cpp index b166ee5a1..79598368c 100644 --- a/common/chat-diff-analyzer.cpp +++ b/common/chat-diff-analyzer.cpp @@ -124,16 +124,16 @@ static std::vector"); analysis.preserved_tokens.push_back(""); analysis.preserved_tokens.push_back(""); analysis.preserved_tokens.push_back(""); diff --git a/common/common.h b/common/common.h index 295451970..8d26337a1 100644 --- a/common/common.h +++ b/common/common.h @@ -1082,6 +1082,9 @@ enum ggml_opt_optimizer_type common_opt_get_optimizer(const char *); struct common_prompt_checkpoint { int64_t n_tokens; + // (optional) id of the task that created the checkpoint + int id_task = -1; + llama_pos pos_min; llama_pos pos_max; diff --git a/embd_res/klite.embd b/embd_res/klite.embd index 3b3f86b40..4b86baea6 100644 --- a/embd_res/klite.embd +++ b/embd_res/klite.embd @@ -7765,7 +7765,7 @@ Current version indicated by LITEVER below. const codeContainer = btn.parentElement.querySelector('pre code'); let innercode = codeContainer.innerText; //remove common language descriptiors from the start - let langsmatched = ["matlab","jsonc","powershell","ps1","haskell","hs","vbnet","vb","apache","apacheconf","makefile","mk","ini","protobuf","proto","typescript","tsx","markdown","md","mkdown","mkd","python","py","javascript","js","jsx","html","xhtml","xml","css","json","typescript","ts","tsx","bash","sh","zsh","java","csharp","cs","c","h","cpp","hpp","php","sql","ruby","rb","go","golang","kotlin","kt","swift","rust","rs","r","dart","scala","dockerfile","docker","yaml","yml","ini","toml","perl","pl","shell","console","powershell","ps1","lua","typescript","ts"]; + let langsmatched = ["matlab","jsonc","powershell","ps1","haskell","hs","vbnet","vb","apache","apacheconf","makefile","mk","ini","protobuf","proto","typescript","tsx","markdown","md","mkdown","mkd","python","py","javascript","js","jsx","html","xhtml","xml","css","json","typescript","ts","tsx","bash","sh","zsh","java","csharp","cs","c","h","cpp","hpp","php","sql","ruby","rb","go","golang","kotlin","kt","swift","rust","rs","r","dart","scala","dockerfile","docker","yaml","yml","ini","toml","perl","pl","shell","console","powershell","ps1","lua","typescript","ts","patch","diff"]; for(let i = 0; i < langsmatched.length; ++i) { let matcher = langsmatched[i]+"\n"; if (innercode.startsWith(matcher)) { diff --git a/ggml/include/ggml-rpc.h b/ggml/include/ggml-rpc.h index 5ad121ae5..efa5420a6 100644 --- a/ggml/include/ggml-rpc.h +++ b/ggml/include/ggml-rpc.h @@ -8,10 +8,10 @@ extern "C" { #define RPC_PROTO_MAJOR_VERSION 4 #define RPC_PROTO_MINOR_VERSION 0 -#define RPC_PROTO_PATCH_VERSION 1 +#define RPC_PROTO_PATCH_VERSION 2 #ifdef __cplusplus -static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); +static_assert(GGML_OP_COUNT == 98, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); #endif #define GGML_RPC_MAX_SERVERS 16 diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index d0ca0676a..f72b0ad2a 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -576,6 +576,7 @@ extern "C" { GGML_OP_RWKV_WKV7, GGML_OP_SOLVE_TRI, GGML_OP_GATED_DELTA_NET, + GGML_OP_LIGHTNING_INDEXER, GGML_OP_UNARY, @@ -2593,6 +2594,24 @@ extern "C" { struct ggml_tensor * state, int64_t K); + // DSA lightning indexer + // + // q: [n_embd_idx, n_head_idx, n_batch, ne3 ] + // k: [n_embd_idx, 1, n_kv, ne3 ] + // weights: [n_head_idx, n_batch, 1, ne3 ] !! prescaled !! + // mask: [n_kv, n_batch, 1, ne33] !! f16 !! + // res: [n_kv, n_batch, 1, ne3 ] + // + // broadcast: + // ne3 % ne33 == 0 + // + GGML_API struct ggml_tensor * ggml_lightning_indexer( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * weights, + struct ggml_tensor * mask); + // custom operators typedef void (*ggml_custom1_op_t)(struct ggml_tensor * dst , const struct ggml_tensor * a, int ith, int nth, void * userdata); diff --git a/ggml/include/gguf.h b/ggml/include/gguf.h index 67851ba6f..b3a1e1230 100644 --- a/ggml/include/gguf.h +++ b/ggml/include/gguf.h @@ -125,12 +125,13 @@ extern "C" { // get ith C string from array with given key_id GGML_API const char * gguf_get_arr_str (const struct gguf_context * ctx, int64_t key_id, size_t i); - GGML_API int64_t gguf_get_n_tensors (const struct gguf_context * ctx); - GGML_API int64_t gguf_find_tensor (const struct gguf_context * ctx, const char * name); // returns -1 if the tensor is not found - GGML_API size_t gguf_get_tensor_offset(const struct gguf_context * ctx, int64_t tensor_id); - GGML_API const char * gguf_get_tensor_name (const struct gguf_context * ctx, int64_t tensor_id); - GGML_API enum ggml_type gguf_get_tensor_type (const struct gguf_context * ctx, int64_t tensor_id); - GGML_API size_t gguf_get_tensor_size (const struct gguf_context * ctx, int64_t tensor_id); + GGML_API int64_t gguf_get_n_tensors (const struct gguf_context * ctx); + GGML_API int64_t gguf_find_tensor (const struct gguf_context * ctx, const char * name); // returns -1 if the tensor is not found + GGML_API size_t gguf_get_tensor_offset(const struct gguf_context * ctx, int64_t tensor_id); + GGML_API const char * gguf_get_tensor_name (const struct gguf_context * ctx, int64_t tensor_id); + GGML_API const int64_t * gguf_get_tensor_ne (const struct gguf_context * ctx, int64_t tensor_id); // returns ne, an array of GGML_MAX_DIMS elements; ne[dim] is 1 for dim >= n_dims + GGML_API enum ggml_type gguf_get_tensor_type (const struct gguf_context * ctx, int64_t tensor_id); + GGML_API size_t gguf_get_tensor_size (const struct gguf_context * ctx, int64_t tensor_id); // removes key if it exists, returns id that the key had prior to removal (-1 if it didn't exist) GGML_API int64_t gguf_remove_key(struct gguf_context * ctx, const char * key); diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 8ab0abf69..80dd75fb5 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -2832,6 +2832,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm { ggml_compute_forward_gated_delta_net(params, tensor); } break; + case GGML_OP_LIGHTNING_INDEXER: + { + ggml_compute_forward_lightning_indexer(params, tensor); + } break; case GGML_OP_MAP_CUSTOM1: { ggml_compute_forward_map_custom1(params, tensor); @@ -3193,6 +3197,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { case GGML_OP_FLASH_ATTN_BACK: case GGML_OP_SSM_CONV: case GGML_OP_SSM_SCAN: + case GGML_OP_LIGHTNING_INDEXER: { n_tasks = n_threads; } break; @@ -3814,6 +3819,12 @@ struct ggml_cplan ggml_graph_plan( // to prepare magnitude and phase for inverted FFTs. cur = ggml_type_size(node->type)*(n_tasks + node->ne[0] * n_tasks * 4); } break; + case GGML_OP_LIGHTNING_INDEXER: + { + // temp buffer for dequantizing lightning indexer keys + const int64_t ne10 = node->src[1]->ne[0]; + cur += sizeof(float)*ne10*n_tasks; + } break; default: break; } diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 0376eb476..d3246f351 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -11568,3 +11568,87 @@ void ggml_compute_forward_fwht(const ggml_compute_params * params, ggml_tensor * } } } + +// ggml_compute_forward_lightning_indexer + +void ggml_compute_forward_lightning_indexer( + const ggml_compute_params * params, + ggml_tensor * dst) { + + const ggml_tensor * q = dst->src[0]; + const ggml_tensor * k = dst->src[1]; + const ggml_tensor * w = dst->src[2]; // weights + const ggml_tensor * m = dst->src[3]; // mask + + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT( q->type == GGML_TYPE_F32); + GGML_ASSERT( w->type == GGML_TYPE_F32); + GGML_ASSERT( m->type == GGML_TYPE_F16); + + GGML_TENSOR_LOCALS(int64_t, neq, q, ne) + GGML_TENSOR_LOCALS(size_t, nbq, q, nb) + GGML_TENSOR_LOCALS(int64_t, nek, k, ne) + GGML_TENSOR_LOCALS(size_t, nbk, k, nb) + GGML_TENSOR_LOCALS(int64_t, new, w, ne) + GGML_TENSOR_LOCALS(size_t, nbw, w, nb) + GGML_TENSOR_LOCALS(int64_t, nem, m, ne) + GGML_TENSOR_LOCALS(size_t, nbm, m, nb) + GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) + GGML_TENSOR_LOCALS(size_t, nb, dst, nb) + + GGML_ASSERT( nb0 == ggml_type_size(dst->type)); + GGML_ASSERT(nbq0 == ggml_type_size( q->type)); + GGML_ASSERT(nbk0 == ggml_type_size( k->type)); + GGML_ASSERT(nbw0 == ggml_type_size( w->type)); + GGML_ASSERT(nbm0 == ggml_type_size( m->type)); + + const int n_embd = q->ne[0]; + const int n_head = q->ne[1]; + const int n_tokens = q->ne[2]; + const int n_stream = q->ne[3]; + const int n_kv = k->ne[2]; + + ggml_to_float_t const k_to_float = ggml_get_type_traits(k->type)->to_float; + GGML_ASSERT((k->type == GGML_TYPE_F32 || k_to_float) && "lightning indexer: unsupported K-type"); + + const int nr = n_kv; + const int ith = params->ith; + const int nth = params->nth; + + // (temporary) buffer for K converted to float + float * k_row_f32 = (float *) params->wdata + ith*(1*n_embd + CACHE_LINE_SIZE_F32); + + // rows per thread + const int dr = (nr + nth - 1)/nth; + + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); + + for (int s = 0; s < n_stream; ++s) { + for (int t = 0; t < n_tokens; ++t) { + const float * w_row = (float *) ((char *) w->data + t*nbw1 + s*nbw3); + const ggml_fp16_t * m_row = (ggml_fp16_t *) ((char *) m->data + t*nbm1 + (s%nem3)*nbm3); + float * dst_row = (float *) ((char *) dst->data + t*nb1 + s*nb3 ); + for (int ik = ir0; ik < ir1; ++ik) { + char * k_row = (char *) k->data + ik*nbk2 + s*nbk3; + if (k_to_float) { + k_to_float(k_row, k_row_f32, n_embd); + } else { + k_row_f32 = (float *) k_row; + } + float score = 0.0f; + for (int h = 0; h < n_head; ++h) { + // dot product of q and k for head h + float qk = 0.0f; + const float * q_row = (float *) ((char *) q->data + h*nbq1 + t*nbq2 + s*nbq3); + ggml_vec_dot_f32(n_embd, &qk, 0, q_row, 0, k_row_f32, 0, 1); + // ReLU and weights (prescaled) + score += MAX(qk, 0.0f) * w_row[h]; + } + // apply mask + dst_row[ik] = score + GGML_CPU_FP16_TO_FP32(m_row[ik]); + } + } + } +} diff --git a/ggml/src/ggml-cpu/ops.h b/ggml/src/ggml-cpu/ops.h index a8e18c716..e956c25d3 100644 --- a/ggml/src/ggml-cpu/ops.h +++ b/ggml/src/ggml-cpu/ops.h @@ -105,6 +105,7 @@ void ggml_compute_forward_rwkv_wkv7(const struct ggml_compute_params * params, s void ggml_compute_forward_solve_tri(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_gla(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_gated_delta_net(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_lightning_indexer(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_map_custom1(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_map_custom2(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_map_custom3(const struct ggml_compute_params * params, struct ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 490bcf885..06c6086de 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -4506,7 +4506,14 @@ static bool ggml_backend_cuda_get_available_uma_memory(long * available_memory_k static void ggml_backend_cuda_device_get_memory(ggml_backend_dev_t dev, size_t * free, size_t * total) { ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context; ggml_cuda_set_device(ctx->device); - CUDA_CHECK(cudaMemGetInfo(free, total)); + cudaError_t err = cudaMemGetInfo(free, total); + if (err != cudaSuccess) { + (void)cudaGetLastError(); + GGML_LOG_WARN("%s: cudaMemGetInfo failed (%s), returning 0/0\n", __func__, cudaGetErrorString(err)); + *free = 0; + *total = 0; + return; + } // ref: https://github.com/ggml-org/llama.cpp/pull/17368 #if defined(__linux__) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index d955ba9ab..872a1a1c6 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -6524,6 +6524,14 @@ static vk_device ggml_vk_get_device(size_t idx) { device->mul_mat_id_m[i] = true; device->mul_mat_id_s[i] = false; break; + case VK_VENDOR_ID_QUALCOMM: + device->mul_mat_l[i] = false; + device->mul_mat_m[i] = true; + device->mul_mat_s[i] = true; + device->mul_mat_id_l[i] = false; + device->mul_mat_id_m[i] = true; + device->mul_mat_id_s[i] = true; + break; #endif default: device->mul_mat_l[i] = true; diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 6ef6e0d1e..a21c96a53 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1095,6 +1095,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "RWKV_WKV7", "SOLVE_TRI", "GATED_DELTA_NET", + "LIGHTNING_INDEXER", "UNARY", @@ -1112,7 +1113,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "GLU", }; -static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT != 97"); +static_assert(GGML_OP_COUNT == 98, "GGML_OP_COUNT != 98"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", @@ -1206,6 +1207,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "rwkv_wkv7(r, w, k, v, a, b, s)", "A X = B, A triangular, solve X", "gated_delta_net(q, k, v, g, beta, s)", + "lightning_indexer(q, k, weights, mask)", "unary(x)", @@ -1223,7 +1225,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "glu(x)", }; -static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT != 97"); +static_assert(GGML_OP_COUNT == 98, "GGML_OP_COUNT != 98"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); @@ -6303,6 +6305,42 @@ struct ggml_tensor * ggml_gated_delta_net( return result; } +// ggml_lightning_indexer + +struct ggml_tensor * ggml_lightning_indexer( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * weights, + struct ggml_tensor * mask) { + + GGML_ASSERT( q->type == GGML_TYPE_F32); + GGML_ASSERT( weights->type == GGML_TYPE_F32); + GGML_ASSERT( mask->type == GGML_TYPE_F16); + GGML_ASSERT( q->ne[0] == k->ne[0]); + GGML_ASSERT( mask->ne[0] == k->ne[2]); + GGML_ASSERT( q->ne[1] == weights->ne[0]); + GGML_ASSERT( k->ne[1] == 1); + GGML_ASSERT( mask->ne[1] == q->ne[2]); + GGML_ASSERT( q->ne[2] == weights->ne[1]); + GGML_ASSERT(weights->ne[2] == 1); + GGML_ASSERT( mask->ne[2] == 1); + GGML_ASSERT( q->ne[3] == k->ne[3]); + GGML_ASSERT( k->ne[3] == weights->ne[3]); + GGML_ASSERT(weights->ne[3] % mask->ne[3] == 0); + + int64_t ne[4] = { k->ne[2], q->ne[2], 1, q->ne[3] }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + result->op = GGML_OP_LIGHTNING_INDEXER; + result->src[0] = q; + result->src[1] = k; + result->src[2] = weights; + result->src[3] = mask; + + return result; +} + //////////////////////////////////////////////////////////////////////////////// struct ggml_hash_set ggml_hash_set_new(size_t size) { diff --git a/ggml/src/gguf.cpp b/ggml/src/gguf.cpp index 04878e9e8..da215eafe 100644 --- a/ggml/src/gguf.cpp +++ b/ggml/src/gguf.cpp @@ -587,6 +587,10 @@ static struct gguf_context * gguf_init_from_reader(const struct gguf_reader & gr GGML_LOG_ERROR("%s: encountered bad_alloc error while reading key %" PRIi64 "\n", __func__, i); ok = false; } + if (ok && key.empty()) { + GGML_LOG_ERROR("%s: key %" PRIi64 " is empty\n", __func__, i); + ok = false; + } for (size_t j = 0; ok && j < ctx->kv.size(); ++j) { if (key == ctx->kv[j].key) { GGML_LOG_ERROR("%s: duplicate key '%s' for tensors %zu and %" PRIi64 " \n", __func__, key.c_str(), j, i); @@ -1226,6 +1230,11 @@ const char * gguf_get_tensor_name(const struct gguf_context * ctx, int64_t tenso return ctx->info[tensor_id].t.name; } +const int64_t * gguf_get_tensor_ne(const struct gguf_context * ctx, int64_t tensor_id) { + GGML_ASSERT(tensor_id >= 0 && tensor_id < gguf_get_n_tensors(ctx)); + return ctx->info[tensor_id].t.ne; +} + enum ggml_type gguf_get_tensor_type(const struct gguf_context * ctx, int64_t tensor_id) { GGML_ASSERT(tensor_id >= 0 && tensor_id < gguf_get_n_tensors(ctx)); return ctx->info[tensor_id].t.type; diff --git a/src/llama-context.cpp b/src/llama-context.cpp index dc120ea15..69a651fd3 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -58,6 +58,12 @@ static const llm_fused_op_probe llm_fused_op_gdn_ch_probe = { /*.n_tokens_per_seq =*/ 16, }; +static const llm_fused_op_probe llm_fused_op_lid_probe = { + /*.op =*/ LLM_FUSED_OP_LIGHTNING_INDEXER, + /*.name =*/ "Lightning Indexer", + /*.n_tokens_per_seq =*/ 1, +}; + llama_context::llama_context( const llama_model & model, llama_context_params params) : @@ -229,6 +235,9 @@ llama_context::llama_context( cparams.fused_gdn_ch = true; cparams.auto_fgdn = true; + cparams.fused_lid = true; + cparams.auto_flid = true; + // with causal attention, the batch size is limited by the context size cparams.n_batch = cparams.causal_attn ? std::min(cparams.n_ctx, params.n_batch) : params.n_batch; @@ -530,6 +539,12 @@ void llama_context::resolve_fused_ops(const llama_memory_context_i * mctx, uint3 resolve(llm_fused_op_gdn_ch_probe, cparams.fused_gdn_ch); cparams.auto_fgdn = false; } + + if (cparams.auto_flid) { + LLAMA_LOG_INFO("%s: resolving fused Lightning Indexer support:\n", func); + resolve(llm_fused_op_lid_probe, cparams.fused_lid); + cparams.auto_flid = false; + } } void llama_context::sched_reserve() { diff --git a/src/llama-cparams.h b/src/llama-cparams.h index 546ae1e2c..58520caa3 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -41,6 +41,8 @@ struct llama_cparams { bool fused_gdn_ar; // use fused gated delta net (autoregressive) bool fused_gdn_ch; // use fused gated delta net (chunked) bool auto_fgdn; + bool fused_lid; // use fused lightning indexer + bool auto_flid; bool no_perf; bool warmup; // TODO: remove [TAG_LLAMA_GRAPH_NO_WARMUP] bool op_offload; diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 6d215a957..f445da9b7 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -843,7 +843,7 @@ static void dsv4_build_comp_inputs( GGML_ASSERT(n_stream > 0); GGML_ASSERT(n_tokens%n_stream == 0); - inp.kq_mask = ggml_new_tensor_4d(ctx, cparams.flash_attn && strcmp(name, "lid") != 0 ? GGML_TYPE_F16 : GGML_TYPE_F32, plan.n_kv, n_tokens/n_stream, 1, n_stream); + inp.kq_mask = ggml_new_tensor_4d(ctx, (strcmp(name, "lid") != 0 && cparams.flash_attn) || (strcmp(name, "lid") == 0 && cparams.fused_lid) ? GGML_TYPE_F16 : GGML_TYPE_F32, plan.n_kv, n_tokens/n_stream, 1, n_stream); ggml_set_input(inp.kq_mask); ggml_set_name(inp.kq_mask, (std::string("dsv4_") + name + "_kq_mask").c_str()); } @@ -3026,9 +3026,9 @@ llm_graph_input_attn_k_dsa * llm_graph_context::build_attn_inp_k_dsa() const { { inp->self_k_idxs_lid = mctx_cur->get_lid()->build_input_k_idxs(ctx0, ubatch); - // ensure F32 mask + // ensure that mask type matches fused lightning indexer use (requires f16 mask) auto cparams_copy = cparams; - cparams_copy.flash_attn = false; + cparams_copy.flash_attn = cparams.fused_lid; inp->self_kq_mask_lid = build_attn_inp_kq_mask(ctx0, mctx_cur->get_lid(), ubatch, cparams_copy); inp->self_kq_mask_lid_cnv = inp->self_kq_mask_lid; diff --git a/src/llama-graph.h b/src/llama-graph.h index 97141ef93..c84cb6a45 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -42,6 +42,7 @@ enum llm_fused_op { LLM_FUSED_OP_FLASH_ATTN, LLM_FUSED_OP_GDN_AR, LLM_FUSED_OP_GDN_CH, + LLM_FUSED_OP_LIGHTNING_INDEXER, }; enum llm_ffn_op_type : int { diff --git a/src/llama-kv-cache-dsv4.cpp b/src/llama-kv-cache-dsv4.cpp index 9fccf347e..ebafac091 100644 --- a/src/llama-kv-cache-dsv4.cpp +++ b/src/llama-kv-cache-dsv4.cpp @@ -29,6 +29,15 @@ static uint32_t dsv4_comp_size(uint32_t kv_size, uint32_t ratio) { return std::max(1, (kv_size + ratio - 1)/ratio); } +static void dsv4_clear_tensor_stream(ggml_tensor * tensor, uint32_t stream) { + GGML_ASSERT(ggml_is_contiguous(tensor)); + GGML_ASSERT(tensor->ne[3] == 1); + GGML_ASSERT(stream < (uint32_t) tensor->ne[2]); + + const size_t stream_size = tensor->nb[2]; + ggml_backend_tensor_memset(tensor, 0, stream*stream_size, stream_size); +} + static int64_t dsv4_stream_offset(uint32_t n_stream, llama_seq_id seq_id, uint32_t size) { if (n_stream <= 1) { return 0; @@ -781,11 +790,20 @@ llama_dsv4_comp_state::llama_dsv4_comp_state( __func__, name, ratio, state_size, n_embd_state, n_stream, layers.size(), total_size()/1024.0/1024.0); } -void llama_dsv4_comp_state::clear(bool data) { +void llama_dsv4_comp_state::clear(llama_seq_id seq_id, bool data) { if (!data) { return; } + if (seq_id >= 0) { + GGML_ASSERT((uint32_t) seq_id < n_stream); + for (const auto & layer : layers) { + dsv4_clear_tensor_stream(layer.kv, (uint32_t) seq_id); + dsv4_clear_tensor_stream(layer.score, (uint32_t) seq_id); + } + return; + } + for (auto & [_, buf] : ctxs_bufs) { ggml_backend_buffer_clear(buf.get(), 0); } @@ -1034,7 +1052,7 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4( // graph does not necessarily overwrite; uninitialized buffer contents would // otherwise leak in (instance-specific garbage) and corrupt recall. Zero all // compressed buffers up front so reads of un-written rows are deterministic. - clear_compressed(true); + clear_compressed(-1, true); } llama_memory_context_ptr llama_kv_cache_dsv4::init_batch( @@ -1147,7 +1165,7 @@ bool llama_kv_cache_dsv4::get_can_shift() const { void llama_kv_cache_dsv4::clear(bool data) { kv_raw->clear(data); - clear_compressed(true); // DSV4 compressed buffers must never expose stale/uninit rows + clear_compressed(-1, true); // DSV4 compressed buffers must never expose stale/uninit rows } bool llama_kv_cache_dsv4::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) { @@ -1169,7 +1187,7 @@ bool llama_kv_cache_dsv4::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1 const bool res = kv_raw->seq_rm(seq_id, p0, p1); if (res) { - clear_compressed(true); + clear_compressed(seq_id, true); } return res; @@ -1177,22 +1195,29 @@ bool llama_kv_cache_dsv4::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1 void llama_kv_cache_dsv4::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) { kv_raw->seq_cp(seq_id_src, seq_id_dst, p0, p1); - clear_compressed(true); } void llama_kv_cache_dsv4::seq_keep(llama_seq_id seq_id) { + GGML_ASSERT(seq_id >= 0 && (uint32_t) seq_id < n_seq_max); + kv_raw->seq_keep(seq_id); - clear_compressed(true); + + for (llama_seq_id id = 0; id < (llama_seq_id) n_seq_max; ++id) { + if (id == seq_id) { + continue; + } + + kv_raw->seq_rm(id, -1, -1); + clear_compressed(id, true); + } } void llama_kv_cache_dsv4::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) { kv_raw->seq_add(seq_id, p0, p1, shift); - clear_compressed(true); } void llama_kv_cache_dsv4::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) { kv_raw->seq_div(seq_id, p0, p1, d); - clear_compressed(true); } llama_pos llama_kv_cache_dsv4::seq_pos_min(llama_seq_id seq_id) const { @@ -1328,13 +1353,32 @@ llama_dsv4_comp_state * llama_kv_cache_dsv4::get_lid_state() const { return lid_state.get(); } -void llama_kv_cache_dsv4::clear_compressed(bool data) { - kv_csa->clear(data); - kv_hca->clear(data); - kv_lid->clear(data); - csa_state->clear(data); - hca_state->clear(data); - lid_state->clear(data); +void llama_kv_cache_dsv4::clear_compressed(llama_seq_id seq_id, bool data) { + if (seq_id < 0) { + kv_csa->clear(data); + kv_hca->clear(data); + kv_lid->clear(data); + } else { + GGML_ASSERT((uint32_t) seq_id < n_seq_max); + + const auto clear_seq = [seq_id, data](llama_kv_cache * kv) { + kv->seq_rm(seq_id, -1, -1); + + if (data) { + for (uint32_t il : kv->get_layer_ids()) { + dsv4_clear_tensor_stream(kv->get_k_storage(il), (uint32_t) seq_id); + } + } + }; + + clear_seq(kv_csa.get()); + clear_seq(kv_hca.get()); + clear_seq(kv_lid.get()); + } + + csa_state->clear(seq_id, data); + hca_state->clear(seq_id, data); + lid_state->clear(seq_id, data); } // diff --git a/src/llama-kv-cache-dsv4.h b/src/llama-kv-cache-dsv4.h index 772b428cd..91619ce9e 100644 --- a/src/llama-kv-cache-dsv4.h +++ b/src/llama-kv-cache-dsv4.h @@ -21,7 +21,7 @@ public: const char * name, const llama_memory_i::layer_filter_cb & filter); - void clear(bool data); + void clear(llama_seq_id seq_id, bool data); uint32_t get_ratio() const; uint32_t get_state_size() const; @@ -67,6 +67,8 @@ private: // DSV4 uses a normal raw/SWA token cache plus compressed K-only block caches. // The compressed caches are storage only; DSV4-specific visibility and block // planning are handled by llama_kv_cache_dsv4_context / llm_graph_input_dsv4. +// FIXME: currently the cache only supports non-unified mode even if unified flag is passed +// FIXME: we currently conflate token_pos and buffer contents. See https://github.com/ggml-org/llama.cpp/pull/25521#discussion_r3558173819 class llama_kv_cache_dsv4 : public llama_memory_i { public: @@ -146,7 +148,7 @@ private: std::unique_ptr hca_state; std::unique_ptr lid_state; - void clear_compressed(bool data); + void clear_compressed(llama_seq_id seq_id, bool data); }; // DSV4 raw attention only uses the SWA half of kv_raw. The base half is kept diff --git a/src/llama-model.cpp b/src/llama-model.cpp index c074bd891..16579c15a 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -451,8 +451,7 @@ llama_model * llama_model_create(llm_arch arch, const llama_model_params & param if (model != nullptr) { model->arch = arch; - auto & devices = model->devices; - if (!devices.empty() && devices[0].is_meta && !llm_arch_supports_sm_tensor(arch)) { + if (params.split_mode == LLAMA_SPLIT_MODE_TENSOR && !llm_arch_supports_sm_tensor(arch)) { throw std::runtime_error(std::string("LLAMA_SPLIT_MODE_TENSOR not implemented for architecture '") + llm_arch_name(arch) + "'"); } } diff --git a/src/models/deepseek32.cpp b/src/models/deepseek32.cpp index 9a20e2ce9..32262e684 100644 --- a/src/models/deepseek32.cpp +++ b/src/models/deepseek32.cpp @@ -301,43 +301,50 @@ llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_ indexer_q = ggml_view_4d(ctx0, indexer_q, indexer_q->ne[0], indexer_q->ne[1], indexer_q->ne[2]/n_stream, n_stream, indexer_q->nb[1], indexer_q->nb[2], indexer_q->nb[3]/n_stream, 0); indexer_weights = ggml_view_4d(ctx0, indexer_weights, indexer_weights->ne[0], indexer_weights->ne[1]/n_stream, indexer_weights->ne[2], n_stream, indexer_weights->nb[1], indexer_weights->nb[2]/n_stream, indexer_weights->nb[3]/n_stream, 0); - // calculate indexer kq - indexer_q = ggml_permute(ctx0, indexer_q, 0, 2, 1, 3); - cb(indexer_q, "indexer_q", il); - indexer_k = ggml_permute(ctx0, indexer_k, 0, 2, 1, 3); - cb(indexer_k, "indexer_k", il); - - ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k, indexer_q); - cb(indexer_kq, "indexer_kq", il); - - // ReLU requires contiguous tensors - indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3)); - cb(indexer_kq, "indexer_kq", il); - - // apply ReLU - ggml_tensor * indexer_score = ggml_relu(ctx0, indexer_kq); - cb(indexer_score, "indexer_score", il); - // pre-scale weights to avoid scaling operations on huge indexer_score tensor indexer_weights = ggml_scale(ctx0, indexer_weights, 1.0f / sqrtf(float(n_embd_indexer_head * n_indexer_head))); cb(indexer_weights, "indexer_weights", il); - // multiply scores by indexer weights - indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights); - cb(indexer_score, "indexer_score", il); + ggml_tensor * indexer_score = nullptr; + if (cparams.fused_lid) { + indexer_score = ggml_lightning_indexer(ctx0, indexer_q, indexer_k, indexer_weights, inp_attn_dsa->get_kq_mask_lid()); + cb(indexer_score, "indexer_score", il); + res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, indexer_score, il}); + } else { + // calculate indexer kq + indexer_q = ggml_permute(ctx0, indexer_q, 0, 2, 1, 3); + cb(indexer_q, "indexer_q", il); + indexer_k = ggml_permute(ctx0, indexer_k, 0, 2, 1, 3); + cb(indexer_k, "indexer_k", il); - // sum by q n_indexer_head dimension - indexer_score = ggml_sum_rows(ctx0, indexer_score); - cb(indexer_score, "indexer_score", il); + ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k, indexer_q); + cb(indexer_kq, "indexer_kq", il); - // permute result to match KQ mask - indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3)); - cb(indexer_score, "indexer_score", il); + // ReLU requires contiguous tensors + indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3)); + cb(indexer_kq, "indexer_kq", il); - // mask indexer scores - ggml_tensor * indexer_kq_mask = inp_attn_dsa->get_kq_mask_lid(); - indexer_score = ggml_add(ctx0, indexer_score, indexer_kq_mask); - cb(indexer_score, "indexer_score", il); + // apply ReLU + indexer_score = ggml_relu(ctx0, indexer_kq); + cb(indexer_score, "indexer_score", il); + + // multiply scores by indexer weights + indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights); + cb(indexer_score, "indexer_score", il); + + // sum by q n_indexer_head dimension + indexer_score = ggml_sum_rows(ctx0, indexer_score); + cb(indexer_score, "indexer_score", il); + + // permute result to match KQ mask + indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3)); + cb(indexer_score, "indexer_score", il); + + // mask indexer scores + ggml_tensor * indexer_kq_mask = inp_attn_dsa->get_kq_mask_lid(); + indexer_score = ggml_add(ctx0, indexer_score, indexer_kq_mask); + cb(indexer_score, "indexer_score", il); + } // get indices of top k indexer scores uint32_t n_top_k = indexer_score->ne[0] < n_indexer_top_k ? indexer_score->ne[0] : n_indexer_top_k; diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index 07aa477e1..e64e43451 100644 --- a/src/models/deepseek4.cpp +++ b/src/models/deepseek4.cpp @@ -556,25 +556,32 @@ ggml_tensor * llama_model_deepseek4::graph::build_lid_top_k( indexer_weights->ne[0], indexer_weights->ne[1]/n_stream, indexer_weights->ne[2], n_stream, indexer_weights->nb[1], indexer_weights->nb[2]/n_stream, indexer_weights->nb[3]/n_stream, 0); - indexer_q = ggml_permute(ctx0, indexer_q, 0, 2, 1, 3); - cb(indexer_q, "lid_q", il); - indexer_k = ggml_permute(ctx0, indexer_k, 0, 2, 1, 3); - cb(indexer_k, "lid_k", il); + ggml_tensor * indexer_score = nullptr; + if (cparams.fused_lid) { + indexer_score = ggml_lightning_indexer(ctx0, indexer_q, indexer_k, indexer_weights, inp_lid.kq_mask); + cb(indexer_score, "lid_score_masked", il); + res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, indexer_score, il}); + } else { + indexer_q = ggml_permute(ctx0, indexer_q, 0, 2, 1, 3); + cb(indexer_q, "lid_q", il); + indexer_k = ggml_permute(ctx0, indexer_k, 0, 2, 1, 3); + cb(indexer_k, "lid_k", il); - ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k, indexer_q); - cb(indexer_kq, "lid_kq", il); + ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k, indexer_q); + cb(indexer_kq, "lid_kq", il); - indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3)); - cb(indexer_kq, "lid_kq", il); + indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3)); + cb(indexer_kq, "lid_kq", il); - ggml_tensor * indexer_score = ggml_relu(ctx0, indexer_kq); - indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights); - indexer_score = ggml_sum_rows(ctx0, indexer_score); - indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3)); - cb(indexer_score, "lid_score", il); + indexer_score = ggml_relu(ctx0, indexer_kq); + indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights); + indexer_score = ggml_sum_rows(ctx0, indexer_score); + indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3)); + cb(indexer_score, "lid_score", il); - indexer_score = ggml_add(ctx0, indexer_score, inp_lid.kq_mask); - cb(indexer_score, "lid_score_masked", il); + indexer_score = ggml_add(ctx0, indexer_score, inp_lid.kq_mask); + cb(indexer_score, "lid_score_masked", il); + } const uint32_t n_top_k = indexer_score->ne[0] < hparams.indexer_top_k ? indexer_score->ne[0] : hparams.indexer_top_k; ggml_tensor * top_k = ggml_cont(ctx0, ggml_top_k(ctx0, indexer_score, n_top_k)); diff --git a/tools/mtmd/mtmd-cli.cpp b/tools/mtmd/mtmd-cli.cpp index 8704ea79d..08288c868 100644 --- a/tools/mtmd/mtmd-cli.cpp +++ b/tools/mtmd/mtmd-cli.cpp @@ -250,7 +250,8 @@ static int eval_message(mtmd_cli_context & ctx, common_chat_msg & msg) { LOG_DBG("formatted_chat.prompt: %s\n", formatted_chat.c_str()); mtmd_input_text text; - text.text = formatted_chat.c_str(); + text.text = formatted_chat.data(); + text.text_len = formatted_chat.size(); text.add_special = add_bos; text.parse_special = true; diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 24b82db21..3d7280a48 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -809,7 +809,7 @@ void mtmd_free(mtmd_context * ctx) { struct mtmd_tokenizer { mtmd_context * ctx; - std::string input_text; + std::string input_text; // note: can contain null bytes; do not use c_str() bool add_special; bool parse_special; const llama_vocab * vocab; @@ -839,9 +839,10 @@ struct mtmd_tokenizer { size_t n_bitmaps) : ctx(ctx) { add_special = text->add_special; parse_special = text->parse_special; - input_text = text->text; vocab = ctx->vocab; + input_text.assign(text->text, text->text_len); + std::vector bitmaps(bmps, bmps + n_bitmaps); auto parts_str = split_text(input_text, ctx->media_marker); size_t i_bm = 0; diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h index 4e6315624..16f303c10 100644 --- a/tools/mtmd/mtmd.h +++ b/tools/mtmd/mtmd.h @@ -67,6 +67,7 @@ struct mtmd_batch; struct mtmd_input_text { const char * text; + size_t text_len; bool add_special; bool parse_special; }; diff --git a/tools/server/README-dev.md b/tools/server/README-dev.md index 882adca09..e81336e5e 100644 --- a/tools/server/README-dev.md +++ b/tools/server/README-dev.md @@ -126,15 +126,15 @@ It is opt in via the `X-Conversation-Id` header on `POST /v1/chat/completions`. The feature lives entirely in `server-stream.{h,cpp}` and rests on three types: -- `stream_session`: a bounded ring buffer (4 MiB cap, oldest bytes drop first) plus a condvar. `append` pushes raw SSE bytes, `read_from` drains from any offset and blocks for live bytes or finalize, `finalize` wakes readers, `cancel` stops the producer. One conv maps to at most one live session. +- `stream_session`: a bounded ring buffer (4 MiB cap, oldest bytes drop first) plus a condvar. `append` pushes raw SSE bytes, `read_from` drains from any offset and blocks for live bytes or finalize, `finalize` wakes readers, `cancel` sets the flag the producer polls. One conv maps to at most one live session. - `stream_session_manager`: a file-static singleton (`g_stream_sessions`) inside `server-stream.cpp`, owns all sessions keyed by conv id, enforces the one conv one session invariant via `create_or_replace`, and runs a GC thread that drops completed sessions past their TTL. Exposed to main only through `server_stream_session_manager_start/stop`. - `stream_pipe_producer` / `stream_pipe_consumer`: the write and read ends. The producer owns the session lifetime and finalizes it on destruction; the consumer is read only and never finalizes, so a reader detaching cannot kill a running generation. -The implementation is hidden in `server-stream.cpp` (pimpl). The header exposes only the route handler factories, `server_stream_session_attach_pipe`, `server_stream_aware_should_stop`, `server_stream_conv_id_from_headers` and the GC lifecycle; the session, manager and consumer types stay in the `.cpp`. +The implementation is hidden in `server-stream.cpp` (pimpl). The header exposes only the route handler factories, the `server_res_spipe` response base, `server_stream_conv_id_from_headers` and the GC lifecycle; the session, manager, consumer and the `server_stream_create_spipe` factory stay in the `.cpp`. -Producer side: `server_res_generator` attaches a producer pipe when the header is present. The HTTP content provider mirrors every chunk into the ring before writing it to the socket. While a pipe is attached, `server_stream_aware_should_stop` ignores peer disconnect, so a dropped socket does not stop generation: only an explicit `DELETE` does. When the peer leaves early, `on_complete` calls `close()`, which drains the rest of the generation into the ring on the http worker. +Producer side: `server_res_generator` extends `server_res_spipe`, which keeps all spipe logic out of the generic `server_http_res`. `set_req` attaches a producer when the header is present, and the wrapped `next` tees each chunk into the ring before the socket, so a chunk lost to a dead wire is already buffered. While attached, `should_stop` ignores peer disconnect: only a `DELETE` stops generation. On an early peer drop, `on_complete` drains the tail into the ring on the http worker. -Lifetime safety: the producer pipe holds a shared `alive` flag also captured by the session cancel hook. `~server_res_generator` calls `cleanup()` to clear that hook while the reader is still alive, so a `cancel` arriving during teardown can never call `stop()` on a freed response. This ordering is the most fragile part of the feature: finalizing or destroying the producer before `cleanup()` runs reintroduces a use after free. +Lifetime safety: the session holds no back reference to the response, so `spipe` is a plain `unique_ptr` touched only by the http worker. `cancel` raises an atomic the producer polls; the producer finalizes the session from its destructor, which also runs `~server_response_reader::stop()` to cancel the generation at the queue level. A `DELETE` stops work by raising the flag and letting the worker unwind. Consumer side: `GET /v1/stream/?from=N` opens a `text/event-stream` that replays buffered bytes from offset `N` and blocks for live bytes, so the browser reattaches like a fresh EventSource. An offset below the dropped prefix returns 400. @@ -235,6 +235,29 @@ That requires `JSON.stringify` when formatted to message content: } ``` +Set `stream: true` in the request body to stream a tool's output as it runs, instead of waiting for it to finish. Only certain tools accept this (for ex. `exec_shell_command`); +returns 404 if tool doesn't support it. + +Response is SSE stream, one `data: ` line per chunk: + +```json +{"chunk": "hello\n"} +``` + +followed by a final event once the tool returns: + +```json +{"done": true} +``` + +or, if `invoke()` threw: + +```json +{"done": true, "error": "..."} +``` + +There is no `[DONE]` sentinel (unlike `/chat/completions`), the stream ends after the `done` + ### Router mode: how child <--> router communicates Upon spawning a new child process using `subprocess`, both child and router listen to the stdout/stderr (combined) diff --git a/tools/server/server-chat.cpp b/tools/server/server-chat.cpp index 02858a2a0..31f94e023 100644 --- a/tools/server/server-chat.cpp +++ b/tools/server/server-chat.cpp @@ -431,22 +431,70 @@ json server_chat_convert_anthropic_to_oai(const json & body) { std::string tool_use_id = json_value(block, "tool_use_id", std::string()); auto result_content = json_value(block, "content", json()); - std::string result_text; if (result_content.is_string()) { - result_text = result_content.get(); + tool_results.push_back({ + {"role", "tool"}, + {"tool_call_id", tool_use_id}, + {"content", result_content.get()} + }); } else if (result_content.is_array()) { + // Single-pass: build both text and content_parts, decide format at the end + std::string result_text; + json content_parts = json::array(); + bool has_images = false; + for (const auto & c : result_content) { - if (json_value(c, "type", std::string()) == "text") { - result_text += json_value(c, "text", std::string()); + std::string c_type = json_value(c, "type", std::string()); + if (c_type == "text") { + std::string text = json_value(c, "text", std::string()); + result_text += text; + content_parts.push_back({ + {"type", "text"}, + {"text", text} + }); + } else if (c_type == "image") { + has_images = true; + json source = json_value(c, "source", json::object()); + std::string source_type = json_value(source, "type", std::string()); + if (source_type == "base64") { + std::string media_type = json_value(source, "media_type", std::string("image/jpeg")); + std::string data = json_value(source, "data", std::string()); + std::string url = "data:" + media_type + ";base64," + data; + content_parts.push_back({ + {"type", "image_url"}, + {"image_url", {{"url", url}}} + }); + } else if (source_type == "url") { + content_parts.push_back({ + {"type", "image_url"}, + {"image_url", {{"url", json_value(source, "url", std::string())}}} + }); + } } } - } - tool_results.push_back({ - {"role", "tool"}, - {"tool_call_id", tool_use_id}, - {"content", result_text} - }); + if (!has_images) { + // Text-only: collapse to a plain string for maximum compatibility + tool_results.push_back({ + {"role", "tool"}, + {"tool_call_id", tool_use_id}, + {"content", result_text} + }); + } else { + // Mixed or image-only: use array content parts (OpenAI multimodal tool format) + tool_results.push_back({ + {"role", "tool"}, + {"tool_call_id", tool_use_id}, + {"content", content_parts} + }); + } + } else { + tool_results.push_back({ + {"role", "tool"}, + {"tool_call_id", tool_use_id}, + {"content", ""} + }); + } } } diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp index ac291d359..78b5c6819 100644 --- a/tools/server/server-common.cpp +++ b/tools/server/server-common.cpp @@ -705,7 +705,8 @@ server_tokens process_mtmd_prompt(mtmd_context * mctx, const std::string & promp std::vector inputs; // multimodal mtmd_input_text inp_txt = { - prompt.c_str(), + prompt.data(), + prompt.size(), /* add_special */ true, /* parse_special */ true, }; @@ -1116,7 +1117,8 @@ json oaicompat_chat_params_parse( // Reasoning budget: pass parameters through to sampling layer { - int reasoning_budget = json_value(body, "thinking_budget_tokens", -1); + int reasoning_budget = json_value(body, "reasoning_budget_tokens", + json_value(body, "thinking_budget_tokens", -1)); if (reasoning_budget == -1) { reasoning_budget = opt.reasoning_budget; } @@ -1125,7 +1127,7 @@ json oaicompat_chat_params_parse( 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_message"] = opt.reasoning_budget_message; + 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-context.cpp b/tools/server/server-context.cpp index 98d0cca1c..5fadc1c8d 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -2290,6 +2290,24 @@ private: // n_tokens_cur: the number of tokens added to the batch for the current slot void create_checkpoint(server_slot & slot, const int64_t n_tokens_cur, llama_pos pos_min, llama_pos pos_max) { + const int id_task = slot.task->id; + + // evict checkpoints within min-step of a previous checkpoint, unless they were + // created by the current task + int64_t last = -1; + for (auto it = slot.prompt.checkpoints.begin(); it != slot.prompt.checkpoints.end(); ) { + if (it->id_task != id_task && last >= 0 && it->n_tokens <= last + params_base.checkpoint_min_step) { + SLT_TRC(slot, "erasing context checkpoint too close to an earlier one (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", size = %.3f MiB)\n", + it->pos_min, it->pos_max, it->n_tokens, (float) it->size() / 1024 / 1024); + + it = slot.prompt.checkpoints.erase(it); + continue; + } + + last = it->n_tokens; + ++it; + } + while (slot.prompt.checkpoints.size() >= (size_t) params_base.n_ctx_checkpoints) { // make room for the new checkpoint, if needed const auto & cur = slot.prompt.checkpoints.front(); @@ -2302,6 +2320,8 @@ private: auto & cur = slot.prompt.checkpoints.emplace_back(); + cur.id_task = id_task; + // [TAG_CHECKPOINTS_FIX_POS_MIN] // TODO: here we incorrectly deterimne that the saved checkpoint data covers the [pos_min, pos_max] range // this is not true for SWA models: https://github.com/ggml-org/llama.cpp/pull/24411#issuecomment-4677983225 @@ -3511,7 +3531,10 @@ private: do_checkpoint = do_checkpoint && !has_mtmd; // no need to create checkpoints that are too close together, unless it's the last user message - do_checkpoint = do_checkpoint && (slot.prompt.checkpoints.empty() || is_last_user_message || n_tokens_start > slot.prompt.checkpoints.back().n_tokens + params_base.checkpoint_min_step); + do_checkpoint = do_checkpoint && ( + slot.prompt.checkpoints.empty() || + is_last_user_message || near_prompt_end || + n_tokens_start > slot.prompt.checkpoints.back().n_tokens + params_base.checkpoint_min_step); SLT_DBG(slot, "main/do_checkpoint = %s, pos_min = %d, pos_max = %d\n", do_checkpoint ? "yes" : "no", pos_min, pos_max); // note: we create the checkpoint before calling llama_decode(), so the current batch is not @@ -3979,11 +4002,9 @@ server_context_meta server_context::get_meta() const { }; } - - // generator-like API for HTTP response generation // may have bypass_sleep = true if the task does not use ctx_server -struct server_res_generator : server_http_res { +struct server_res_generator : server_res_spipe { server_response_reader rd; server_res_generator(server_queue & queue_tasks, server_response & queue_results, int sleep_idle_seconds, bool bypass_sleep = false) : rd(queue_tasks, queue_results, HTTP_POLLING_SECONDS) { @@ -3993,15 +4014,6 @@ struct server_res_generator : server_http_res { queue_tasks.wait_until_no_sleep(); } } - ~server_res_generator() override { - // cleanup() must run while rd is still alive (rd is destroyed after this body returns) - if (spipe) { - spipe->cleanup(); - } - } - void stop() override { - rd.stop(); - } void ok(const json & response_data) { status = 200; data = safe_json_to_str(response_data); @@ -4039,6 +4051,8 @@ std::unique_ptr server_routes::handle_completions_impl( auto & rd = res->rd; auto & params = this->params; + res->set_req(&req); // will also set spipe if needed + int32_t sse_ping_interval = params.sse_ping_interval; try { @@ -4181,7 +4195,7 @@ std::unique_ptr server_routes::handle_completions_impl( } res->status = 200; res->content_type = "text/event-stream"; - res->next = [res_this = res.get(), res_type, sse_ping_interval, &req](std::string & output) -> bool { + res->set_next([res_this = res.get(), res_type, sse_ping_interval](std::string & output) -> bool { static auto format_error = [](task_response_type res_type, const json & res_json) { if (res_type == TASK_RESPONSE_TYPE_ANTHROPIC) { return format_anthropic_sse({ @@ -4193,7 +4207,9 @@ std::unique_ptr server_routes::handle_completions_impl( } }; - auto effective_should_stop = server_stream_aware_should_stop(res_this, req.should_stop); + auto effective_should_stop = [&res_this]() { + return res_this->should_stop(); + }; try { if (effective_should_stop()) { @@ -4284,13 +4300,9 @@ std::unique_ptr server_routes::handle_completions_impl( // terminate on exception return false; } - }; + }); } - // attach a producer pipe to the response when X-Conversation-Id is present. - // the pipe mirrors SSE chunks into the ring buffer and wires up the cancel hook. - server_stream_session_attach_pipe(*res, req.headers); - return res; } diff --git a/tools/server/server-http.cpp b/tools/server/server-http.cpp index 87eee5fd4..48f903dfc 100644 --- a/tools/server/server-http.cpp +++ b/tools/server/server-http.cpp @@ -1,7 +1,6 @@ #include "common.h" #include "http.h" #include "server-http.h" -#include "server-stream.h" #include "server-common.h" #include "ui.h" @@ -530,33 +529,20 @@ static void process_handler_response(server_http_req_ptr && request, server_http std::string chunk; const bool has_next = response->next(chunk); if (!chunk.empty()) { - // mirror into the ring buffer first, the session must reflect every SSE chunk - // whether or not the wire write below succeeds - if (response->spipe) { - response->spipe->write(chunk.data(), chunk.size()); - } if (!sink.write(chunk.data(), chunk.size())) { - // peer is gone, stop the wire path here return false; } SRV_DBG("http: streamed chunk: %s\n", chunk.c_str()); } if (!has_next) { - // producer reached its natural end on the wire, a later close() skips the drain - if (response->spipe) { - response->spipe->done(); - } sink.done(); SRV_DBG("%s", "http: stream ended\n"); } return has_next; }; const auto on_complete = [request = q_ptr, response = r_ptr](bool) mutable { - // on a dropped peer, close() drains the rest of the generation into the ring buffer - if (response->spipe) { - response->spipe->close(); - } - response.reset(); // spipe destructor finalizes the session if attached + response->on_complete(); + response.reset(); request.reset(); }; res.set_chunked_content_provider(content_type, chunked_content_provider, on_complete); @@ -564,6 +550,7 @@ static void process_handler_response(server_http_req_ptr && request, server_http res.status = response->status; set_headers(res, response->headers); res.set_content(response->data, response->content_type); + response->on_complete(); } } diff --git a/tools/server/server-http.h b/tools/server/server-http.h index 350813183..032b08d0d 100644 --- a/tools/server/server-http.h +++ b/tools/server/server-http.h @@ -11,7 +11,6 @@ #include struct common_params; -struct stream_pipe_producer; // defined in server-stream.h // generator-like API for HTTP response generation // this object response with one of the 2 modes: @@ -25,19 +24,13 @@ struct server_http_res { std::string data; std::map headers; - // if set, the stream survives a client disconnect: the producer pipe keeps draining into the - // ring buffer and finalizes the session on destruction, so no explicit on_stream_end is needed. - // shared_ptr (not unique_ptr) so the forward-declared type is safe to delete here. - std::shared_ptr spipe; - std::function next = nullptr; bool is_stream() const { return next != nullptr; } - // called when the session is cancelled (e.g. DELETE /v1/stream/). - // server_res_generator overrides this to stop its reader; the default is a no-op. - virtual void stop() {} + // fired before req and res are destroyed + virtual void on_complete() {} virtual ~server_http_res() = default; }; diff --git a/tools/server/server-models.cpp b/tools/server/server-models.cpp index d1fdc0607..7c6bfda0f 100644 --- a/tools/server/server-models.cpp +++ b/tools/server/server-models.cpp @@ -219,13 +219,14 @@ void server_model_meta::update_caps() { "LLAMA_ARG_MODEL_URL", "LLAMA_ARG_MMPROJ", "LLAMA_ARG_MMPROJ_URL", + "LLAMA_ARG_MMPROJ_AUTO", "LLAMA_ARG_HF_REPO", "LLAMA_ARG_HF_REPO_FILE", }); params.offline = true; common_models_handler handler = common_models_handler_init(params, LLAMA_EXAMPLE_SERVER); common_models_handler_apply(handler, params); // note: this won't download the model because offline=true - if (params.mmproj.path.empty()) { + if (params.no_mmproj || params.mmproj.path.empty()) { multimodal = { false, false }; } else { multimodal = mtmd_get_cap_from_file(params.mmproj.path.c_str()); diff --git a/tools/server/server-stream.cpp b/tools/server/server-stream.cpp index 553ac26b1..19db04988 100644 --- a/tools/server/server-stream.cpp +++ b/tools/server/server-stream.cpp @@ -96,8 +96,6 @@ struct stream_session { size_t dropped_prefix() const; // bytes evicted from the front due to cap int64_t completed_at() const; // 0 while alive, unix seconds after finalize - void set_stop_producer(std::function fn); - void cancel(); private: @@ -109,7 +107,6 @@ private: bool done; std::atomic cancelled; // polled lock-free by the should_stop closure, no mu int64_t completed_ts; - std::function stop_producer; }; stream_session::stream_session(std::string conversation_id_, size_t max_bytes_) : conversation_id(std::move(conversation_id_)) @@ -217,26 +214,10 @@ int64_t stream_session::completed_at() const { return completed_ts; } -void stream_session::set_stop_producer(std::function fn) { - std::lock_guard lock(mu); - stop_producer = std::move(fn); -} - void stream_session::cancel() { - // flip cancelled first so the producer-side server_stream_aware_should_stop can break out of the - // recv() wait even if remove_waiting_task_ids does not notify the condvar (the cancel task - // posted by rd.stop() will eventually notify, but we do not want to depend on that timing) + // the should_stop closure on both the producer and any HTTP reader polls is_cancelled() + // so flipping this is the only signal needed to unwind both sides cancelled.store(true, std::memory_order_release); - // copy the hook under the lock then invoke outside, the producer side may grab queue locks - // and we do not want to hold our mu across that path - std::function fn; - { - std::lock_guard lock(mu); - fn = stop_producer; - } - if (fn) { - fn(); - } } bool stream_session::is_cancelled() const { @@ -325,8 +306,10 @@ void stream_session_manager::evict_and_cancel(const std::string & conversation_i s = it->second; sessions.erase(it); } - // signal the producer side first so the inference is cancelled at the queue level, - // then finalize, which wakes any pending HTTP reader and lets the drain exit naturally + // cancel first so the producer's on_complete() drain loop and any pending HTTP reader + // observe is_cancelled() and stop pulling further output, then finalize to wake readers + // blocked in read_from(). note: this does not interrupt the underlying generation itself, + // which keeps running to its own natural stop condition (EOS/max_tokens) s->cancel(); s->finalize(); } @@ -431,65 +414,15 @@ stream_pipe_producer::stream_pipe_producer(stream_session_ptr session) } stream_pipe_producer::~stream_pipe_producer() { - cleanup(); session_->finalize(); } -void stream_pipe_producer::cleanup() { - if (!alive_) { - return; - } - alive_->store(false, std::memory_order_release); - session_->set_stop_producer(nullptr); - alive_.reset(); -} - bool stream_pipe_producer::write(const char * data, size_t len) { return session_->append(data, len); } -void stream_pipe_producer::done() { - done_ = true; -} - -void stream_pipe_producer::close() { - // httplib bails its content provider the moment is_peer_alive() goes false, so pump the rest - // of the generation into the ring buffer here. a DELETE flips is_cancelled and cuts it short - if (done_ || session_->is_cancelled()) { - SRV_TRC("stream_pipe close: skip drain (done=%d cancelled=%d) conv=%s\n", - done_ ? 1 : 0, session_->is_cancelled() ? 1 : 0, session_->conversation_id.c_str()); - return; - } - SRV_TRC("stream_pipe close: draining conv=%s\n", session_->conversation_id.c_str()); - size_t drained = 0; - std::string chunk; - while (true) { - chunk.clear(); - bool has_next = res_->next(chunk); - if (!chunk.empty()) { - write(chunk.data(), chunk.size()); - drained += chunk.size(); - } - if (!has_next) { - break; - } - } - SRV_TRC("stream_pipe close: drain ended conv=%s bytes=%zu\n", session_->conversation_id.c_str(), drained); -} - -std::shared_ptr stream_pipe_producer::create(stream_session_ptr session, - server_http_res & res) { - auto alive = std::make_shared>(true); - auto * res_ptr = &res; - session->set_stop_producer([alive, res_ptr]() { - if (alive->load(std::memory_order_acquire)) { - res_ptr->stop(); - } - }); - auto pipe = std::shared_ptr(new stream_pipe_producer(std::move(session))); - pipe->alive_ = std::move(alive); - pipe->res_ = res_ptr; - return pipe; +stream_pipe_producer * stream_pipe_producer::create(stream_session_ptr session) { + return new stream_pipe_producer(std::move(session)); } // stream_pipe_consumer @@ -661,21 +594,68 @@ std::string server_stream_conv_id_from_headers(const std::map & headers) { +static stream_pipe_producer * server_stream_create_spipe(const std::map & headers) { std::string conversation_id = server_stream_conv_id_from_headers(headers); SRV_TRC("conv_id=%s (empty=%d)\n", conversation_id.c_str(), conversation_id.empty() ? 1 : 0); if (conversation_id.empty()) { - return; + return nullptr; } auto session = g_stream_sessions.create_or_replace(conversation_id); - res.spipe = stream_pipe_producer::create(session, res); + return stream_pipe_producer::create(session); } -std::function server_stream_aware_should_stop(server_http_res * res, std::function fallback) { - return [res, fallback = std::move(fallback)]() -> bool { - if (res->spipe) { - return res->spipe->is_cancelled(); +// +// server_res_spipe +// + +void server_res_spipe::set_req(const server_http_req * req) { + this->req = req; + // optionally attach spipe to the response when X-Conversation-Id is present + spipe.reset(server_stream_create_spipe(req->headers)); +} + +bool server_res_spipe::conn_alive() { + GGML_ASSERT(req != nullptr); + return !req->should_stop(); +} + +bool server_res_spipe::should_stop() { + if (spipe) { + // note: if DELETE /v1/stream/ is called, is_cancelled() will be true + return spipe->is_cancelled(); + } else { + return !conn_alive(); + } +} + +void server_res_spipe::on_complete() { + if (!spipe || next_finished) { + return; + } + std::string chunk; + while (!spipe->is_cancelled()) { + chunk.clear(); + bool has_next = next_orig(chunk); + if (!chunk.empty()) { + spipe->write(chunk.data(), chunk.size()); } - return fallback(); + if (!has_next) { + break; + } + } +} + +void server_res_spipe::set_next(std::function next_fn) { + next_orig = std::move(next_fn); + next = [this](std::string & out) { + bool has_next = next_orig(out); + if (spipe) { + // if spipe is set, tee-style pipe input to both HTTP and spipe + spipe->write(out.data(), out.size()); + } + if (!has_next) { + next_finished = true; + } + return has_next; }; } diff --git a/tools/server/server-stream.h b/tools/server/server-stream.h index c0c3e924f..9753140dd 100644 --- a/tools/server/server-stream.h +++ b/tools/server/server-stream.h @@ -30,36 +30,15 @@ protected: // producer end: writes chunks into the ring buffer and owns the session lifetime, finalizing it // on destruction. -// -// lifetime safety: holds a shared_ptr> alive also captured by the session's -// stop_producer hook. cleanup() sets alive=false and clears the hook; it must run while the -// response the hook calls stop() on is still alive. ~server_res_generator() does this explicitly. struct stream_pipe_producer : stream_pipe { ~stream_pipe_producer() override; bool write(const char * data, size_t len); - // mark the natural end on the wire so a later close() is a no-op - void done(); - - // on a peer drop, pump the response next() into the ring buffer until done. runs on the http - // worker from on_complete, no-op after done() or cancel - void close(); - - // disarm the stop hook and drop the alive guard, must run while the response the hook - // references is still alive. idempotent, the destructor calls it too - void cleanup(); - - // res.stop() is invoked when the session is cancelled, the alive guard ensures stop() is not - // called after cleanup() has run - static std::shared_ptr create(stream_session_ptr session, server_http_res & res); + static stream_pipe_producer * create(stream_session_ptr session); private: explicit stream_pipe_producer(stream_session_ptr session); - - bool done_ = false; - std::shared_ptr> alive_; - server_http_res * res_ = nullptr; }; void server_stream_session_manager_start(); @@ -73,10 +52,22 @@ server_http_context::handler_t server_stream_make_delete_handler(); // extract the X-Conversation-Id header value (case-insensitive), empty when absent std::string server_stream_conv_id_from_headers(const std::map & headers); -// on an X-Conversation-Id header, create or replace the session and attach a producer pipe to res -void server_stream_session_attach_pipe(server_http_res & res, const std::map & headers); +// implement tee-style pipe (spipe) for "stream replay" functionality +struct server_res_spipe : server_http_res { +private: + // if set, the stream survives a client disconnect: + // connection kept alive, output is forwarded to spipe and reuse later + std::unique_ptr spipe; + // if spipe is set, use this next_orig to implement tee-style pipe + std::function next_orig; + const server_http_req * req = nullptr; + // set once next_orig reports no more data, so on_complete() doesn't re-drain a finished stream + bool next_finished = false; -// should_stop closure that ignores peer disconnect when a pipe is attached, so only an explicit -// DELETE stops the producer and generation keeps flowing into the ring buffer. without a pipe it -// delegates to fallback, the legacy non-resumable flow -std::function server_stream_aware_should_stop(server_http_res * res, std::function fallback); +public: + void set_req(const server_http_req * req); + bool conn_alive(); + bool should_stop(); + void on_complete() override; + void set_next(std::function next_fn); +}; diff --git a/tools/server/server-tools.cpp b/tools/server/server-tools.cpp index 3f162a13e..a8216d7db 100644 --- a/tools/server/server-tools.cpp +++ b/tools/server/server-tools.cpp @@ -12,6 +12,7 @@ #include #include #include +#include namespace fs = std::filesystem; @@ -51,7 +52,13 @@ public: virtual bool write_file(const std::string & path, const std::string & content) const = 0; // paths relative to `base`, '/'-separated; sets `err` if `base` isn't a directory virtual std::vector list_files(const std::string & base, std::string & err) const = 0; - virtual exec_result run(const std::vector & args, size_t max_output, int timeout_secs) const = 0; + // on_chunk, if set, is called with each chunk of output as it is read (before truncation cuts in); + // returning false terminates the process early (e.g. the client disconnected) + virtual exec_result run( + const std::vector & args, + size_t max_output, + int timeout_secs, + const std::function & on_chunk = nullptr) const = 0; }; class tools_io_basic : public tools_io { @@ -123,7 +130,11 @@ public: return list_files_fallback(base); } - exec_result run(const std::vector & args, size_t max_output, int timeout_secs) const override { + exec_result run( + const std::vector & args, + size_t max_output, + int timeout_secs, + const std::function & on_chunk = nullptr) const override { exec_result res; subprocess_s proc; @@ -164,8 +175,14 @@ public: size_t len = strlen(buf); if (output.size() + len <= max_output) { output.append(buf, len); + if (on_chunk && !on_chunk(std::string(buf, len))) { + subprocess_terminate(&proc); + break; + } } else { - output.append(buf, max_output - output.size()); + size_t remaining = max_output - output.size(); + output.append(buf, remaining); + if (on_chunk && remaining > 0) on_chunk(std::string(buf, remaining)); truncated = true; } } @@ -287,7 +304,7 @@ struct server_tool_read_file : server_tool { }; } - json invoke(json params) const override { + json invoke(json params, server_tool::stream *) const override { std::string path = params.at("path").get(); int start_line = json_value(params, "start_line", 1); int end_line = json_value(params, "end_line", -1); // -1 = no limit @@ -376,7 +393,7 @@ struct server_tool_file_glob_search : server_tool { }; } - json invoke(json params) const override { + json invoke(json params, server_tool::stream *) const override { std::string base = params.at("path").get(); std::string include = json_value(params, "include", std::string("**")); std::string exclude = json_value(params, "exclude", std::string("")); @@ -457,7 +474,7 @@ struct server_tool_grep_search : server_tool { }; } - json invoke(json params) const override { + json invoke(json params, server_tool::stream *) const override { std::string path = params.at("path").get(); std::string pat_str = params.at("pattern").get(); std::string include = json_value(params, "include", std::string("**")); @@ -577,6 +594,7 @@ struct server_tool_exec_shell_command : server_tool { name = "exec_shell_command"; display_name = "Execute shell command"; permission_write = true; + support_stream = true; } json get_definition() const override { @@ -598,7 +616,7 @@ struct server_tool_exec_shell_command : server_tool { }; } - json invoke(json params) const override { + json invoke(json params, server_tool::stream * st) const override { std::string command = params.at("command").get(); int timeout = json_value(params, "timeout", 10); size_t max_output = (size_t) json_value(params, "max_output_size", (int) SERVER_TOOL_EXEC_SHELL_COMMAND_MAX_OUTPUT_SIZE); @@ -612,7 +630,24 @@ struct server_tool_exec_shell_command : server_tool { std::vector args = {"sh", "-c", command}; #endif - auto io = make_tools_io(params); + auto io = make_tools_io(params); + + if (st) { + auto res = io->run(args, max_output, timeout, [st](const std::string & chunk) { + st->push(chunk); + return !st->alive || st->alive(); + }); + if (st->alive && !st->alive()) { + return json(); + } + std::string tail = string_format("\n[exit code: %d]", res.exit_code); + if (res.timed_out) { + tail += " [exit due to timed out]"; + } + st->push(tail); + return json(); + } + auto res = io->run(args, max_output, timeout); std::string text_output = res.output; @@ -654,7 +689,7 @@ struct server_tool_write_file : server_tool { }; } - json invoke(json params) const override { + json invoke(json params, server_tool::stream *) const override { std::string path = params.at("path").get(); std::string content = params.at("content").get(); @@ -710,7 +745,7 @@ struct server_tool_edit_file : server_tool { }; } - json invoke(json params) const override { + json invoke(json params, server_tool::stream *) const override { std::string path = params.at("path").get(); const json & edits_json = params.at("edits"); @@ -1018,7 +1053,7 @@ struct server_tool_get_datetime : server_tool { }; } - json invoke(json) const override { + 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); @@ -1026,6 +1061,59 @@ struct server_tool_get_datetime : server_tool { } }; +struct server_tool_stream_result : server_task_result { + std::string chunk; + bool done = false; + std::string error_msg; + + json to_json() override { + if (!done) { + return {{"chunk", chunk}}; + } else { + json result = {{"done", true}}; + if (!error_msg.empty()) { + result["error"] = error_msg; + } + return result; + } + } +}; + +void server_tool::stream::push(const std::string & chunk) { + if (chunk.empty()) return; + auto r = std::make_unique(); + r->id = id; + r->chunk = chunk; + qr.send(std::move(r)); +} + +struct server_tools_res : server_http_res { + std::thread worker; + server_response * qr = nullptr; // set only for streaming responses + int id = -1; + + ~server_tools_res() override { + if (worker.joinable()) { + worker.join(); + } + if (qr) { + qr->remove_waiting_task_id(id); + } + } +}; + +static server_tool & find_tool(std::vector> & tools, const std::string & name, bool require_stream) { + for (auto & t : tools) { + if (t->name == name) { + if (require_stream && !t->support_stream) { + throw std::invalid_argument(string_format("tool \"%s\" does not support stream = true", name.c_str())); + } + return *t; + } + } + throw std::invalid_argument(string_format("unknown tool \"%s\"", name.c_str())); +} + // // public API // @@ -1090,16 +1178,63 @@ void server_tools::setup(const std::vector & enabled_tools) { }; handle_post = [this](const server_http_req & req) -> server_http_res_ptr { - auto res = std::make_unique(); + auto res = std::make_unique(); try { json body = json::parse(req.body); std::string tool_name = body.at("tool").get(); json params = body.value("params", json::object()); - json result = invoke(tool_name, params); - res->data = safe_json_to_str(result); + bool stream = body.value("stream", false); + + server_tool & tool = find_tool(tools, tool_name, stream); + + if (stream) { + int id = res_id.fetch_add(1); + queue_res.add_waiting_task_id(id); + res->qr = &queue_res; + res->id = id; + + res->worker = std::thread([this, id, &req, &tool, params]() mutable { + server_tool::stream st{queue_res, id, [&req]() { + return !req.should_stop(); + }}; + + auto done = std::make_unique(); + try { + tool.invoke(params, &st); + } catch (const std::exception & e) { + done->error_msg = e.what(); + } catch (...) { + done->error_msg = "An unknown error occurred"; + } + done->id = st.id; + done->done = true; + st.qr.send(std::move(done)); + }); + + res->content_type = "text/event-stream"; + res->status = 200; + res->next = [this, id](std::string & output) -> bool { + auto result = queue_res.recv(id); + auto * r = dynamic_cast(result.get()); + GGML_ASSERT(r != nullptr); + output = "data: " + safe_json_to_str(r->to_json()) + "\n\n"; + if (r->done) { + queue_res.remove_waiting_task_id(id); + return false; + } + return true; + }; + } else { + json result = tool.invoke(params, nullptr); + res->status = 200; + res->data = safe_json_to_str(result); + } } catch (const json::exception & e) { res->status = 400; res->data = safe_json_to_str(format_error_response(e.what(), ERROR_TYPE_INVALID_REQUEST)); + } catch (const std::invalid_argument & e) { + res->status = 404; + res->data = safe_json_to_str(format_error_response(e.what(), ERROR_TYPE_INVALID_REQUEST)); } catch (const std::exception & e) { SRV_ERR("got exception: %s\n", e.what()); res->status = 500; @@ -1108,12 +1243,3 @@ void server_tools::setup(const std::vector & enabled_tools) { return res; }; } - -json server_tools::invoke(const std::string & name, const json & params) { - for (auto & t : tools) { - if (t->name == name) { - return t->invoke(params); - } - } - return {{"error", "unknown tool: " + name}}; -} diff --git a/tools/server/server-tools.h b/tools/server/server-tools.h index cc147379b..6f6528f48 100644 --- a/tools/server/server-tools.h +++ b/tools/server/server-tools.h @@ -2,15 +2,27 @@ #include "server-common.h" #include "server-http.h" +#include "server-queue.h" + +#include +#include struct server_tool { std::string name; std::string display_name; bool permission_write = false; + bool support_stream = false; // if true, output can be streamed virtual ~server_tool() = default; virtual json get_definition() const = 0; - virtual json invoke(json params) const = 0; + + struct stream { + server_response & qr; + int id; + std::function alive; + void push(const std::string & chunk); + }; + virtual json invoke(json params, stream * st = nullptr) const = 0; json to_json() const; }; @@ -18,8 +30,11 @@ struct server_tool { struct server_tools { std::vector> tools; + // for streaming + server_response queue_res; + std::atomic res_id{0}; + void setup(const std::vector & enabled_tools); - json invoke(const std::string & name, const json & params); server_http_context::handler_t handle_get; server_http_context::handler_t handle_post; diff --git a/tools/server/tests/unit/test_compat_anthropic.py b/tools/server/tests/unit/test_compat_anthropic.py index ef1948d4a..e23947cdd 100644 --- a/tools/server/tests/unit/test_compat_anthropic.py +++ b/tools/server/tests/unit/test_compat_anthropic.py @@ -402,6 +402,65 @@ def test_anthropic_tool_result_with_text(): assert len(res.body["content"]) > 0 +def test_anthropic_tool_result_with_image(): + """Test tool result containing mixed text and image blocks + + Verifies that image blocks inside Anthropic tool_result content are + properly converted to OpenAI image_url format rather than being + silently dropped. With a non-multimodal model, the converted image + triggers a clear error message instead of being ignored. + """ + server.jinja = True + server.start() + + # Small 1x1 red PNG image in base64 (same as vision tests) + red_pixel_png = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg==" + + res = server.make_request("POST", "/v1/messages", data={ + "model": "test", + "max_tokens": 100, + "messages": [ + {"role": "user", "content": "What is in this image?"}, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "tool_1", + "name": "read", + "input": {"file": "test.png"} + } + ] + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "tool_1", + "content": [ + {"type": "text", "text": "File: test.png"}, + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": red_pixel_png + } + } + ] + } + ] + } + ] + }) + + # Without the fix, image block would cause "unsupported content[].type" + # With the fix, image is converted to image_url but tinyllama doesn't support images + assert res.status_code == 500 + assert "image input is not supported" in res.body.get("error", {}).get("message", "").lower() + + def test_anthropic_tool_result_error(): """Test tool result with error flag""" server.jinja = True diff --git a/tools/server/tests/unit/test_tools_builtin.py b/tools/server/tests/unit/test_tools_builtin.py index d4fd5dc9b..1b2d0db43 100755 --- a/tools/server/tests/unit/test_tools_builtin.py +++ b/tools/server/tests/unit/test_tools_builtin.py @@ -105,6 +105,24 @@ def test_tools_builtin_edit_file_rejects_non_unique_old_text(): os.remove(log_path) +def test_tools_builtin_exec_shell_command_stream(): + global server + server.start() + + events = list(server.make_stream_request("POST", "/tools", data={ + "tool": "exec_shell_command", + "params": {"command": "echo hello"}, + "stream": True, + })) + + assert len(events) >= 2 + assert events[-1]["done"] is True + assert not events[-1].get("error") + chunks = "".join(e["chunk"] for e in events[:-1]) + assert "hello" in chunks + assert "[exit code: 0]" in chunks + + def test_tools_builtin_edit_file_rejects_overlapping_edits(): global server server.start() diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddMcpServersSubmenu.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddMcpServersSubmenu.svelte index a75f45f37..f86222823 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddMcpServersSubmenu.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddMcpServersSubmenu.svelte @@ -17,7 +17,7 @@ let { onMcpSettingsClick }: Props = $props(); let mcpSearchQuery = $state(''); - let allMcpServers = $derived(mcpStore.getServersSorted()); + let allMcpServers = $derived(mcpStore.getServers()); let mcpServers = $derived(mcpStore.visibleMcpServers); let hasMcpServers = $derived(mcpServers.length > 0); // let hasAnyMcpServers = $derived(allMcpServers.length > 0); diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddToolsSubmenu.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddToolsSubmenu.svelte index 0ffdf2deb..5f8ff8a11 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddToolsSubmenu.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddToolsSubmenu.svelte @@ -10,7 +10,7 @@ import { useToolsPanel } from '$lib/hooks/use-tools-panel.svelte'; const toolsPanel = useToolsPanel(); - const hasMcpServersAvailable = $derived(mcpStore.getServersSorted().length > 0); + const hasMcpServersAvailable = $derived(mcpStore.getServers().length > 0); open && toolsPanel.handleOpen()}> diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormPickerMcpPrompts/ChatFormPickerMcpPrompts.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormPickerMcpPrompts/ChatFormPickerMcpPrompts.svelte index ff734ac88..f35d816de 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormPickerMcpPrompts/ChatFormPickerMcpPrompts.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormPickerMcpPrompts/ChatFormPickerMcpPrompts.svelte @@ -322,7 +322,7 @@ } let filteredPrompts = $derived.by(() => { - const sortedServers = mcpStore.getServersSorted(); + const sortedServers = mcpStore.getServers(); const serverOrderMap = new Map(sortedServers.map((server, index) => [server.id, index])); const sortedPrompts = [...prompts].sort((a, b) => { diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormPickerMcpResources.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormPickerMcpResources.svelte index 1125ae8ec..ed97e1fc7 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormPickerMcpResources.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormPickerMcpResources.svelte @@ -138,7 +138,7 @@ } let filteredResources = $derived.by(() => { - const sortedServers = mcpStore.getServersSorted(); + const sortedServers = mcpStore.getServers(); const serverOrderMap = new Map(sortedServers.map((server, index) => [server.id, index])); const sortedResources = [...resources].sort((a, b) => { diff --git a/tools/ui/src/lib/components/app/dialogs/DialogMcpServerRecommendations.svelte b/tools/ui/src/lib/components/app/dialogs/DialogMcpServerRecommendations.svelte deleted file mode 100644 index cdbc055ee..000000000 --- a/tools/ui/src/lib/components/app/dialogs/DialogMcpServerRecommendations.svelte +++ /dev/null @@ -1,210 +0,0 @@ - - - - - - Do more with MCP - - Power-up your experience by adding tools, resources and more capabilities provided by MCP - servers. - - - -
-

Quickly get started with

- - {#each RECOMMENDED_MCP_SERVERS as server (server.id)} - (selected[server.id] = enabled)} - /> - {/each} - - {#if addedServers.length > 0} - {#each addedServers as server (server.id)} - - {/each} - {/if} - - {#if showAddForm} - - (newServerUrl = v)} - onHeadersChange={(v) => (newServerHeaders = v)} - urlError={newServerUrl ? newServerUrlError : null} - id="recommendation-new-server" - /> - -
- - - -
- - {:else} - - - - {/if} -
- - - - - - - - diff --git a/tools/ui/src/lib/components/app/dialogs/index.ts b/tools/ui/src/lib/components/app/dialogs/index.ts index 73f22c565..29136308c 100644 --- a/tools/ui/src/lib/components/app/dialogs/index.ts +++ b/tools/ui/src/lib/components/app/dialogs/index.ts @@ -18,15 +18,6 @@ */ export { default as DialogMcpServerAddNew } from './DialogMcpServerAddNew.svelte'; -/** - * **DialogMcpServerRecommendations** - Suggested MCP servers opt-in dialog - * - * Prompts the user to enable pre-defined recommended MCP servers on first launch. - * Shows one switch per suggested server and persists the choice as a per-chat - * override so the selected servers become available in conversations. - */ -export { default as DialogMcpServerRecommendations } from './DialogMcpServerRecommendations.svelte'; - /** * **DialogExportSettings** - Settings export dialog with sensitive data warning * diff --git a/tools/ui/src/lib/components/app/mcp/McpActiveServersAvatars.svelte b/tools/ui/src/lib/components/app/mcp/McpActiveServersAvatars.svelte index 2f732cfd5..dd239423e 100644 --- a/tools/ui/src/lib/components/app/mcp/McpActiveServersAvatars.svelte +++ b/tools/ui/src/lib/components/app/mcp/McpActiveServersAvatars.svelte @@ -13,7 +13,7 @@ let { class: className = '', onclick }: Props = $props(); - let mcpServers = $derived(mcpStore.getServersSorted().filter((s) => s.enabled)); + let mcpServers = $derived(mcpStore.getServers().filter((s) => s.enabled)); let enabledMcpServersForChat = $derived( mcpServers.filter((s) => conversationsStore.isMcpServerEnabledForChat(s.id) && s.url.trim()) ); diff --git a/tools/ui/src/lib/components/app/mcp/McpServerCard/McpServerCardCompact.svelte b/tools/ui/src/lib/components/app/mcp/McpServerCard/McpServerCardCompact.svelte deleted file mode 100644 index 6cb3e18b6..000000000 --- a/tools/ui/src/lib/components/app/mcp/McpServerCard/McpServerCardCompact.svelte +++ /dev/null @@ -1,156 +0,0 @@ - - - -
-
- {#if showSkeleton} - - - - - {:else} - - {/if} -
- - -
- - {#if isError && errorMessage} -

{errorMessage}

- {/if} - - {#if showSkeleton} -
- -
- -
- - - - -
- {:else} - {#if description} - {#if description.lines === 2} -

- {description.text} -

- {:else} -

- {description.text} -

- {/if} - {/if} - - {#if tools.length > 0} -
- {#each visibleTools as tool (tool.name)} - - - - {tool.name} - - - - -

- {tool.description ?? 'No description'} -

- - - {/each} - - {#if hiddenToolCount > 0} - - - - + {hiddenToolCount} more tools - - - - -

- {hiddenTools.map((tool) => tool.name).join(', ')} -

- - - {/if} -
- {/if} - {/if} - diff --git a/tools/ui/src/lib/components/app/mcp/index.ts b/tools/ui/src/lib/components/app/mcp/index.ts index b142841f4..3d30bb3b4 100644 --- a/tools/ui/src/lib/components/app/mcp/index.ts +++ b/tools/ui/src/lib/components/app/mcp/index.ts @@ -180,16 +180,6 @@ export { default as McpServerCardDeleteDialog } from './McpServerCard/McpServerC /** Skeleton loading state for server card during health checks. */ export { default as McpServerCardSkeleton } from './McpServerCardSkeleton.svelte'; -/** - * **McpServerCardCompact** - Condensed MCP server card - * - * Compact alternative to McpServerCard tailored for picker-style UIs. - * Shows the server identity, status, and a flex-wrapped list of available tools. - * Tool names are rendered as badges; hovering a badge shows its description in a tooltip. - * Does not show connection logs or server instructions. - */ -export { default as McpServerCardCompact } from './McpServerCard/McpServerCardCompact.svelte'; - /** * **McpServerIdentity** - Server identity display (icon, name, version) * diff --git a/tools/ui/src/lib/components/app/settings/SettingsMcpServers.svelte b/tools/ui/src/lib/components/app/settings/SettingsMcpServers.svelte index e35cfadbb..120e7914f 100644 --- a/tools/ui/src/lib/components/app/settings/SettingsMcpServers.svelte +++ b/tools/ui/src/lib/components/app/settings/SettingsMcpServers.svelte @@ -1,9 +1,10 @@ -
+
@@ -87,53 +77,78 @@

MCP Servers

- -
-
- {#if servers.length === 0 && !isAddingServer} -
- No MCP Servers configured yet. Add one to enable agentic features. -
- {/if} + {#if servers.length === 0} +
+ + + + + - {#if servers.length > 0} -
- {#each servers as server (server.id)} - {#if !initialLoadComplete} - - {:else} - { - const wasEnabled = conversationsStore.isMcpServerEnabledForChat(server.id); - await conversationsStore.toggleMcpServerForChat(server.id); - if (!wasEnabled) { - toolsStore.enableAllToolsForServer(server.id); - } - }} - onUpdate={(updates) => mcpStore.updateServer(server.id, updates)} - onDelete={() => mcpStore.removeServer(server.id)} - /> - {/if} - {/each} -
- {/if} -
+ Add your first MCP server + + Connect a remote MCP server by URL. + + + + + + +
+ {:else} +
+ {#each servers as server (server.id)} + {#if isServerPending(server.id)} + + {:else} + { + const wasEnabled = conversationsStore.isMcpServerEnabledForChat(server.id); + await conversationsStore.toggleMcpServerForChat(server.id); + if (!wasEnabled) { + toolsStore.enableAllToolsForServer(server.id); + } + }} + onUpdate={(updates) => mcpStore.updateServer(server.id, updates)} + onDelete={() => mcpStore.removeServer(server.id)} + /> + {/if} + {/each} + + {#if !isAddingServer} + + + + + + + Add another MCP server + + Connect a remote MCP server by URL. + + + + + + + {/if} +
+ {/if} diff --git a/tools/ui/src/lib/components/ui/empty/empty-content.svelte b/tools/ui/src/lib/components/ui/empty/empty-content.svelte new file mode 100644 index 000000000..cbae3ab04 --- /dev/null +++ b/tools/ui/src/lib/components/ui/empty/empty-content.svelte @@ -0,0 +1,23 @@ + + +
+ {@render children?.()} +
diff --git a/tools/ui/src/lib/components/ui/empty/empty-description.svelte b/tools/ui/src/lib/components/ui/empty/empty-description.svelte new file mode 100644 index 000000000..4d0fd7d53 --- /dev/null +++ b/tools/ui/src/lib/components/ui/empty/empty-description.svelte @@ -0,0 +1,23 @@ + + +
a:hover]:text-primary text-sm/relaxed [&>a]:underline [&>a]:underline-offset-4', + className + )} + {...restProps} +> + {@render children?.()} +
diff --git a/tools/ui/src/lib/components/ui/empty/empty-header.svelte b/tools/ui/src/lib/components/ui/empty/empty-header.svelte new file mode 100644 index 000000000..87014feaf --- /dev/null +++ b/tools/ui/src/lib/components/ui/empty/empty-header.svelte @@ -0,0 +1,20 @@ + + +
+ {@render children?.()} +
diff --git a/tools/ui/src/lib/components/ui/empty/empty-media.svelte b/tools/ui/src/lib/components/ui/empty/empty-media.svelte new file mode 100644 index 000000000..13e15918c --- /dev/null +++ b/tools/ui/src/lib/components/ui/empty/empty-media.svelte @@ -0,0 +1,41 @@ + + + + +
+ {@render children?.()} +
diff --git a/tools/ui/src/lib/components/ui/empty/empty-title.svelte b/tools/ui/src/lib/components/ui/empty/empty-title.svelte new file mode 100644 index 000000000..83c9810eb --- /dev/null +++ b/tools/ui/src/lib/components/ui/empty/empty-title.svelte @@ -0,0 +1,20 @@ + + +
+ {@render children?.()} +
diff --git a/tools/ui/src/lib/components/ui/empty/empty.svelte b/tools/ui/src/lib/components/ui/empty/empty.svelte new file mode 100644 index 000000000..6c38c10a9 --- /dev/null +++ b/tools/ui/src/lib/components/ui/empty/empty.svelte @@ -0,0 +1,23 @@ + + +
+ {@render children?.()} +
diff --git a/tools/ui/src/lib/components/ui/empty/index.ts b/tools/ui/src/lib/components/ui/empty/index.ts new file mode 100644 index 000000000..cae5ff914 --- /dev/null +++ b/tools/ui/src/lib/components/ui/empty/index.ts @@ -0,0 +1,22 @@ +import Root from './empty.svelte'; +import Header from './empty-header.svelte'; +import Media from './empty-media.svelte'; +import Title from './empty-title.svelte'; +import Description from './empty-description.svelte'; +import Content from './empty-content.svelte'; + +export { + Root, + Header, + Media, + Title, + Description, + Content, + // + Root as Empty, + Header as EmptyHeader, + Media as EmptyMedia, + Title as EmptyTitle, + Description as EmptyDescription, + Content as EmptyContent +}; diff --git a/tools/ui/src/lib/constants/index.ts b/tools/ui/src/lib/constants/index.ts index 643864ad4..b982a5907 100644 --- a/tools/ui/src/lib/constants/index.ts +++ b/tools/ui/src/lib/constants/index.ts @@ -8,7 +8,6 @@ export * from './attachment-labels'; export * from './database'; export * from './reasoning-effort'; export * from './reasoning-effort-tokens'; -export * from './recommended-mcp-servers'; export * from './storage'; export * from './attachment-menu'; export * from './auto-scroll'; diff --git a/tools/ui/src/lib/constants/mcp-form.ts b/tools/ui/src/lib/constants/mcp-form.ts index dbbdc6f12..7a1ccffb0 100644 --- a/tools/ui/src/lib/constants/mcp-form.ts +++ b/tools/ui/src/lib/constants/mcp-form.ts @@ -1,4 +1,2 @@ export const MCP_SERVER_URL_PLACEHOLDER = 'https://mcp.example.com/sse'; export const MIN_AUTOCOMPLETE_INPUT_LENGTH = 1; -/** Number of tools shown on the compact MCP server card before collapsing to a "+ N more" badge */ -export const MCP_CARD_VISIBLE_TOOL_LIMIT = 4; diff --git a/tools/ui/src/lib/constants/recommended-mcp-servers.ts b/tools/ui/src/lib/constants/recommended-mcp-servers.ts deleted file mode 100644 index 1fb3ba8c7..000000000 --- a/tools/ui/src/lib/constants/recommended-mcp-servers.ts +++ /dev/null @@ -1,35 +0,0 @@ -import { DEFAULT_MCP_CONFIG } from './mcp'; -import type { RecommendedMCPServer } from '$lib/types'; - -/** - * Pre-defined recommended MCP servers. - * - * Servers are enabled by default, but they are not turned on for individual - * conversations until the user explicitly enables them (so their tools are - * disabled by default). - */ -export const RECOMMENDED_MCP_SERVERS: RecommendedMCPServer[] = [ - { - id: 'exa-web-search', - name: 'Exa Web Search', - description: 'Search the web and retrieve relevant content.', - url: 'https://mcp.exa.ai/mcp', - enabled: true, - requestTimeoutSeconds: DEFAULT_MCP_CONFIG.requestTimeoutSeconds - }, - { - id: 'huggingface-mcp', - name: 'Hugging Face', - description: - 'Browse models, datasets, spaces and machine learning papers from the Hugging Face hub.', - url: 'https://huggingface.co/mcp', - enabled: true, - requestTimeoutSeconds: DEFAULT_MCP_CONFIG.requestTimeoutSeconds - } -]; - -export const RECOMMENDED_MCP_SERVER_IDS = new Set( - RECOMMENDED_MCP_SERVERS.map((server) => server.id) -); - -export const RECOMMENDED_MCP_SERVERS_OPTIN_DIALOG_DELAY = 1000; diff --git a/tools/ui/src/lib/constants/settings-keys.ts b/tools/ui/src/lib/constants/settings-keys.ts index ea8963044..c69c47f64 100644 --- a/tools/ui/src/lib/constants/settings-keys.ts +++ b/tools/ui/src/lib/constants/settings-keys.ts @@ -58,7 +58,6 @@ export const SETTINGS_KEYS = { // MCP MCP_SERVERS: 'mcpServers', MCP_REQUEST_TIMEOUT_SECONDS: 'mcpRequestTimeoutSeconds', - MCP_DEFAULT_SERVER_OVERRIDES: 'mcpDefaultServerOverrides', AGENTIC_MAX_TURNS: 'agenticMaxTurns', AGENTIC_MAX_TOOL_PREVIEW_LINES: 'agenticMaxToolPreviewLines', SHOW_TOOL_CALL_IN_PROGRESS: 'showToolCallInProgress', diff --git a/tools/ui/src/lib/constants/settings-registry.ts b/tools/ui/src/lib/constants/settings-registry.ts index 347a08dfe..0e04996f2 100644 --- a/tools/ui/src/lib/constants/settings-registry.ts +++ b/tools/ui/src/lib/constants/settings-registry.ts @@ -28,7 +28,6 @@ import McpLogo from '$lib/components/app/mcp/McpLogo.svelte'; import { SETTINGS_KEYS } from './settings-keys'; import { ROUTES, SETTINGS_SECTION_SLUGS } from './routes'; import { TITLE_GENERATION } from './title-generation'; -import { RECOMMENDED_MCP_SERVERS } from './recommended-mcp-servers'; export const SETTINGS_SECTION_TITLES = { GENERAL: 'General', @@ -775,16 +774,9 @@ const NON_UI_SETTINGS: SettingsEntry[] = [ key: SETTINGS_KEYS.MCP_SERVERS, label: 'MCP servers', help: 'Configure MCP servers as a JSON list. Use the form in the MCP Client settings section to edit.', - defaultValue: JSON.stringify(RECOMMENDED_MCP_SERVERS), + defaultValue: '[]', type: SettingsFieldType.INPUT, sync: { serverKey: SETTINGS_KEYS.MCP_SERVERS, paramType: SyncableParameterType.STRING } - }, - { - key: SETTINGS_KEYS.MCP_DEFAULT_SERVER_OVERRIDES, - label: 'MCP default server overrides', - help: 'Per-server enable/disable defaults inherited by new chats. JSON-serialized list of {serverId, enabled} entries.', - defaultValue: '[]', - type: SettingsFieldType.INPUT } // { // key: SETTINGS_KEYS.PY_INTERPRETER_ENABLED, diff --git a/tools/ui/src/lib/constants/storage.ts b/tools/ui/src/lib/constants/storage.ts index eca9739ba..4f130be94 100644 --- a/tools/ui/src/lib/constants/storage.ts +++ b/tools/ui/src/lib/constants/storage.ts @@ -22,8 +22,6 @@ export const DISABLED_TOOLS_LOCALSTORAGE_KEY = `${STORAGE_APP_NAME}.disabledTool export const DISABLED_TOOL_KEYS_LOCALSTORAGE_KEY = `${STORAGE_APP_NAME}.disabledToolKeys`; export const FAVORITE_MODELS_LOCALSTORAGE_KEY = `${STORAGE_APP_NAME}.favoriteModels`; export const REASONING_EFFORT_DEFAULT_LOCALSTORAGE_KEY = `${STORAGE_APP_NAME}.reasoningEffortDefault`; -/** Set when user has interacted with the MCP server recommendations dialog (checked servers, added custom server, or dismissed) */ -export const MCP_SERVERS_ADDED_TO_CHAT_LOCALSTORAGE_KEY = `${STORAGE_APP_NAME}.mcpServersSetupDone`; export const USER_OVERRIDES_LOCALSTORAGE_KEY = `${STORAGE_APP_NAME}.userOverrides`; /** Key prefix for per-conversation resumable stream state, conversationId is appended */ diff --git a/tools/ui/src/lib/hooks/use-mcp-recommendations.svelte.ts b/tools/ui/src/lib/hooks/use-mcp-recommendations.svelte.ts deleted file mode 100644 index 4f4c2c782..000000000 --- a/tools/ui/src/lib/hooks/use-mcp-recommendations.svelte.ts +++ /dev/null @@ -1,80 +0,0 @@ -import { browser } from '$app/environment'; -import { - MCP_SERVERS_ADDED_TO_CHAT_LOCALSTORAGE_KEY, - RECOMMENDED_MCP_SERVER_IDS, - RECOMMENDED_MCP_SERVERS_OPTIN_DIALOG_DELAY -} from '$lib/constants'; -import { mcpStore } from '$lib/stores/mcp.svelte'; - -/** - * First-run opt-in dialog for the recommended MCP servers. - * - * Owns the dismissed / open / trigger-timeout state and the effect that - * schedules the dialog. Reads opt-in status and the configured server list - * from `mcpStore`, so callers don't need to recompute on their side. - */ -export function useMcpRecommendations() { - let dismissed = $state( - browser && localStorage.getItem(MCP_SERVERS_ADDED_TO_CHAT_LOCALSTORAGE_KEY) === 'true' - ); - let open = $state(false); - let checked = $state(false); - let triggerTimeout: ReturnType | null = null; - - function dismiss() { - if (browser) { - localStorage.setItem(MCP_SERVERS_ADDED_TO_CHAT_LOCALSTORAGE_KEY, 'true'); - } - dismissed = true; - open = false; - if (triggerTimeout) { - clearTimeout(triggerTimeout); - triggerTimeout = null; - } - } - - function handleOpenChange(next: boolean) { - open = next; - if (!next) dismiss(); - } - - $effect(() => { - if (!browser) return; - - if (open || dismissed) { - if (triggerTimeout) { - clearTimeout(triggerTimeout); - triggerTimeout = null; - } - return; - } - - // Already evaluated once this session; leave any pending trigger alone so - // it can still fire later. Setting `checked = true` below re-runs this - // effect, and we must not wipe the timeout that was just scheduled. - if (checked) return; - - const hasRecommendations = mcpStore - .getServers() - .some((server) => RECOMMENDED_MCP_SERVER_IDS.has(server.id)); - - if (hasRecommendations) { - triggerTimeout = setTimeout(() => { - open = true; - }, RECOMMENDED_MCP_SERVERS_OPTIN_DIALOG_DELAY); - } - - checked = true; - }); - - return { - get open() { - return open; - }, - get dismissed() { - return dismissed; - }, - dismiss, - handleOpenChange - }; -} diff --git a/tools/ui/src/lib/hooks/use-tools-panel.svelte.ts b/tools/ui/src/lib/hooks/use-tools-panel.svelte.ts index 9e65e2f81..59d736731 100644 --- a/tools/ui/src/lib/hooks/use-tools-panel.svelte.ts +++ b/tools/ui/src/lib/hooks/use-tools-panel.svelte.ts @@ -95,7 +95,7 @@ export function useToolsPanel(): UseToolsPanelReturn { if (toolsStore.builtinTools.length === 0 && !toolsStore.loading) { toolsStore.fetchBuiltinTools(); } - mcpStore.runHealthChecksForServers(mcpStore.getServersSorted().filter((s) => s.enabled)); + mcpStore.runHealthChecksForServers(mcpStore.getServers().filter((s) => s.enabled)); } return { diff --git a/tools/ui/src/lib/services/migration.service.ts b/tools/ui/src/lib/services/migration.service.ts index 981283be9..c5e4093f3 100644 --- a/tools/ui/src/lib/services/migration.service.ts +++ b/tools/ui/src/lib/services/migration.service.ts @@ -522,7 +522,7 @@ const mcpDefaultEnabledMigration: Migration = { const config = configRaw ? JSON.parse(configRaw) : {}; // Don't overwrite an existing config entry — current data wins. - if (SETTINGS_KEYS.MCP_DEFAULT_SERVER_OVERRIDES in config) { + if (MCP_DEFAULT_OVERRIDES_LEGACY_KEY in config) { if (import.meta.env.DEV && import.meta.env.VITE_DEBUG) console.log('[Migration] MCP default enabled: config already has overrides, skipping'); return; @@ -543,7 +543,7 @@ const mcpDefaultEnabledMigration: Migration = { return; } - config[SETTINGS_KEYS.MCP_DEFAULT_SERVER_OVERRIDES] = raw; + config[MCP_DEFAULT_OVERRIDES_LEGACY_KEY] = raw; localStorage.setItem(CONFIG_LOCALSTORAGE_KEY, JSON.stringify(config)); if (import.meta.env.DEV && import.meta.env.VITE_DEBUG) @@ -586,6 +586,83 @@ const configTypesMigration: Migration = { } }; +const MCP_DEFAULT_OVERRIDES_LEGACY_KEY = `${STORAGE_APP_NAME}.mcpDefaultServerOverrides`; +const MCP_DEFAULT_OVERRIDES_MERGE_MIGRATION_ID = 'mcp-default-overrides-merge-v1'; + +/** + * Folds `mcpDefaultServerOverrides` (the legacy "default for new chats" list, + * JSON-encoded as `[{ serverId, enabled }, ...]`) into `mcpServers[i].enabled`. + * The legacy override key is intentionally left in the config so a downgrade + * keeps reading it. Runs after `mcpDefaultEnabledMigration` so any legacy + * standalone overrides are already inside the config. + */ +const mcpDefaultOverridesMergeMigration: Migration = { + id: MCP_DEFAULT_OVERRIDES_MERGE_MIGRATION_ID, + description: + 'Merge mcpDefaultServerOverrides entries onto mcpServers[i].enabled (preserves legacy key)', + + async run(): Promise { + const configRaw = localStorage.getItem(CONFIG_LOCALSTORAGE_KEY); + if (configRaw === null) return; + + const config = JSON.parse(configRaw); + const raw = config[MCP_DEFAULT_OVERRIDES_LEGACY_KEY]; + + if (typeof raw !== 'string' || raw.length === 0) { + if (import.meta.env.DEV && import.meta.env.VITE_DEBUG) + console.log('[Migration] MCP default overrides merge: nothing to merge'); + return; + } + + let overrides: { serverId: string; enabled: boolean }[]; + try { + const parsed = JSON.parse(raw); + if (!Array.isArray(parsed)) return; + overrides = parsed.filter( + (o) => + typeof o === 'object' && + o !== null && + typeof (o as Record).serverId === 'string' && + typeof (o as Record).enabled === 'boolean' + ) as { serverId: string; enabled: boolean }[]; + } catch { + return; + } + + const serversRaw = config[SETTINGS_KEYS.MCP_SERVERS]; + let servers: { id: string; enabled?: boolean }[]; + try { + servers = typeof serversRaw === 'string' ? JSON.parse(serversRaw) : []; + } catch { + return; + } + + if (!Array.isArray(servers)) servers = []; + + let serversChanged = false; + const knownIds = new Set(servers.map((s) => s.id)); + for (const override of overrides) { + if (!knownIds.has(override.serverId)) continue; + const index = servers.findIndex((s) => s.id === override.serverId); + + if (index >= 0 && servers[index].enabled !== override.enabled) { + servers[index] = { ...servers[index], enabled: override.enabled }; + serversChanged = true; + } + } + + if (serversChanged) { + config[SETTINGS_KEYS.MCP_SERVERS] = JSON.stringify(servers); + localStorage.setItem(CONFIG_LOCALSTORAGE_KEY, JSON.stringify(config)); + } + + if (import.meta.env.DEV && import.meta.env.VITE_DEBUG) + console.log( + `[Migration] MCP default overrides merge: applied=${overrides.length} serversChanged=${serversChanged} (legacy key preserved)` + ); + } +}; + const migrations: Migration[] = [ localStorageMigration, idxdbMigration, @@ -593,6 +670,7 @@ const migrations: Migration[] = [ themeMigration, customJsonKeyMigration, mcpDefaultEnabledMigration, + mcpDefaultOverridesMergeMigration, configTypesMigration ]; diff --git a/tools/ui/src/lib/stores/conversations.svelte.ts b/tools/ui/src/lib/stores/conversations.svelte.ts index 47d2f060d..ea2c11e14 100644 --- a/tools/ui/src/lib/stores/conversations.svelte.ts +++ b/tools/ui/src/lib/stores/conversations.svelte.ts @@ -23,7 +23,8 @@ import { browser } from '$app/environment'; import { toast } from 'svelte-sonner'; import { DatabaseService } from '$lib/services/database.service'; import { MigrationService } from '$lib/services/migration.service'; -import { config, settingsStore } from '$lib/stores/settings.svelte'; +import { config } from '$lib/stores/settings.svelte'; +import { mcpStore } from '$lib/stores/mcp.svelte'; import { filterByLeafNodeId, findLeafNode, generateConversationTitle } from '$lib/utils'; import type { McpServerOverride } from '$lib/types/database'; import { zipSync, unzipSync, strToU8, strFromU8 } from 'fflate'; @@ -46,7 +47,6 @@ import { ISO_TIME_SEPARATOR_REPLACEMENT, NON_ALPHANUMERIC_REGEX, MULTIPLE_UNDERSCORE_REGEX, - SETTINGS_KEYS, REASONING_EFFORT_DEFAULT_LOCALSTORAGE_KEY } from '$lib/constants'; @@ -80,9 +80,6 @@ class ConversationsStore { /** Whether the store has been initialized */ isInitialized = $state(false); - /** Pending MCP server overrides for new conversations (before first message) */ - pendingMcpServerOverrides = $state(ConversationsStore.loadMcpDefaults()); - /** Global (non-conversation-specific) thinking toggle default, derived from reasoning effort */ pendingThinkingEnabled = $state(false); @@ -94,28 +91,6 @@ class ConversationsStore { /** Last non-off reasoning effort, restored when re-enabling thinking globally */ private lastNonOffEffort: ReasoningEffort | null = null; - private static loadMcpDefaults(): McpServerOverride[] { - const raw = config()[SETTINGS_KEYS.MCP_DEFAULT_SERVER_OVERRIDES]; - if (typeof raw !== 'string' || raw.length === 0) return []; - try { - const parsed = JSON.parse(raw); - if (!Array.isArray(parsed)) return []; - return parsed.filter( - (o: unknown) => typeof o === 'object' && o !== null && 'serverId' in o && 'enabled' in o - ) as McpServerOverride[]; - } catch { - return []; - } - } - - private saveMcpDefaults(): void { - const plain = this.pendingMcpServerOverrides.map((o) => ({ - serverId: o.serverId, - enabled: o.enabled - })); - settingsStore.updateConfig(SETTINGS_KEYS.MCP_DEFAULT_SERVER_OVERRIDES, JSON.stringify(plain)); - } - /** Load reasoning effort default from localStorage */ private static loadReasoningEffortDefault(): ReasoningEffort | ReasoningEffort.OFF { if (typeof globalThis.localStorage === 'undefined') return ReasoningEffort.OFF; @@ -162,11 +137,6 @@ class ConversationsStore { try { await MigrationService.runAllMigrations(); - - // Re-read defaults after migrations: a migration may have populated - // the settings config (e.g. moved legacy MCP overrides into it). - this.pendingMcpServerOverrides = ConversationsStore.loadMcpDefaults(); - await this.loadConversations(); this.isInitialized = true; } catch (error) { @@ -273,18 +243,9 @@ class ConversationsStore { const conversationName = name || `Chat ${new Date().toLocaleString()}`; const conversation = await DatabaseService.createConversation(conversationName); - if (this.pendingMcpServerOverrides.length > 0) { - // Deep clone to plain objects (Svelte 5 $state uses Proxies which can't be cloned to IndexedDB) - const plainOverrides = this.pendingMcpServerOverrides.map((o) => ({ - serverId: o.serverId, - enabled: o.enabled - })); - conversation.mcpServerOverrides = plainOverrides; - await DatabaseService.updateConversation(conversation.id, { - mcpServerOverrides: plainOverrides - }); - this.pendingMcpServerOverrides = []; - } + // New conversations inherit per-server enabled defaults directly from + // `mcpServers[i].enabled` (see #checkServerEnabled). No per-conversation + // override list needs to be seeded. // Inherit global thinking/reasoning defaults into the new conversation const thinkingEnabled = this.getThinkingEnabled(); @@ -321,7 +282,6 @@ class ConversationsStore { return false; } - this.pendingMcpServerOverrides = []; this.activeConversation = conversation; if (conversation.currNode) { @@ -351,7 +311,6 @@ class ConversationsStore { this.activeConversation = null; this.activeMessages = []; // reload defaults so new chats inherit persisted state - this.pendingMcpServerOverrides = ConversationsStore.loadMcpDefaults(); this.pendingReasoningEffort = ConversationsStore.loadReasoningEffortDefault(); } @@ -641,11 +600,30 @@ class ConversationsStore { * */ + /** + /** + * Resolve the per-server enabled value when no active conversation exists. + * The default for new chats is the server's own `enabled` flag in `mcpServers`. + */ + #getDefaultOverrideForNoConversation(serverId: string): McpServerOverride | undefined { + const server = mcpStore.getServers().find((s) => s.id === serverId); + if (!server) return undefined; + return { serverId, enabled: server.enabled }; + } + + /** + * Default overrides for new chats are derived from `mcpServers[i].enabled`, + * so the global on/off state lives in one place. + */ + #getAllDefaultOverridesForNoConversation(): McpServerOverride[] { + return mcpStore.getServers().map((s) => ({ serverId: s.id, enabled: s.enabled })); + } + /** * Gets MCP server override for a specific server in the active conversation. - * Falls back to pending overrides if no active conversation exists. + * Falls back to `mcpServers[i].enabled` if no active conversation exists. * @param serverId - The server ID to check - * @returns The override if set, undefined if using global setting + * @returns The override if set, undefined if no matching server */ getMcpServerOverride(serverId: string): McpServerOverride | undefined { if (this.activeConversation) { @@ -653,18 +631,18 @@ class ConversationsStore { (o: McpServerOverride) => o.serverId === serverId ); } - return this.pendingMcpServerOverrides.find((o) => o.serverId === serverId); + return this.#getDefaultOverrideForNoConversation(serverId); } /** * Get all MCP server overrides for the current conversation. - * Returns pending overrides if no active conversation. + * When no active conversation, derives from `mcpServers[i].enabled`. */ getAllMcpServerOverrides(): McpServerOverride[] { if (this.activeConversation?.mcpServerOverrides) { return this.activeConversation.mcpServerOverrides; } - return this.pendingMcpServerOverrides; + return this.#getAllDefaultOverridesForNoConversation(); } /** @@ -679,13 +657,16 @@ class ConversationsStore { /** * Sets or removes MCP server override for the active conversation. - * If no conversation exists, stores as pending override. + * If no conversation exists, persists `enabled` onto `mcpServers[i].enabled` + * (the single source of truth for new-chat defaults). * @param serverId - The server ID to override - * @param enabled - The enabled state, or undefined to remove override + * @param enabled - The enabled state, or undefined to remove per-conversation override */ async setMcpServerOverride(serverId: string, enabled: boolean | undefined): Promise { if (!this.activeConversation) { - this.setPendingMcpServerOverride(serverId, enabled); + if (enabled !== undefined) { + mcpStore.updateServer(serverId, { enabled }); + } return; } @@ -729,29 +710,6 @@ class ConversationsStore { } } - /** - * Sets or removes a pending MCP server override (for new conversations). - */ - private setPendingMcpServerOverride(serverId: string, enabled: boolean | undefined): void { - if (enabled === undefined) { - this.pendingMcpServerOverrides = this.pendingMcpServerOverrides.filter( - (o) => o.serverId !== serverId - ); - } else { - const existingIndex = this.pendingMcpServerOverrides.findIndex( - (o) => o.serverId === serverId - ); - if (existingIndex >= 0) { - const newOverrides = [...this.pendingMcpServerOverrides]; - newOverrides[existingIndex] = { serverId, enabled }; - this.pendingMcpServerOverrides = newOverrides; - } else { - this.pendingMcpServerOverrides = [...this.pendingMcpServerOverrides, { serverId, enabled }]; - } - } - this.saveMcpDefaults(); - } - /** * Toggles MCP server enabled state for the active conversation. * @param serverId - The server ID to toggle @@ -769,14 +727,6 @@ class ConversationsStore { await this.setMcpServerOverride(serverId, undefined); } - /** - * Clears all pending MCP server overrides. - */ - clearPendingMcpServerOverrides(): void { - this.pendingMcpServerOverrides = []; - this.saveMcpDefaults(); - } - /** * Gets the effective thinking-enabled state for the active conversation. * Returns the conversation override if set, otherwise the global default. diff --git a/tools/ui/src/lib/stores/mcp.svelte.ts b/tools/ui/src/lib/stores/mcp.svelte.ts index a53463ee9..c752e7a36 100644 --- a/tools/ui/src/lib/stores/mcp.svelte.ts +++ b/tools/ui/src/lib/stores/mcp.svelte.ts @@ -470,18 +470,12 @@ class MCPStore { } } - // Fallback: try favicon from root domain - const fallbackUrl = this.#getServerFaviconFallback(server.url); - if (fallbackUrl) { - return fallbackUrl; - } - - return null; + return this.#getServerFaviconFallback(server.url); } /** * Construct a fallback favicon URL from the MCP server URL. - * e.g. https://mcp.exa.ai/mcp -> https://exa.ai/favicon.ico + * e.g. https://mcp.example.com/sse -> https://example.com/favicon.ico */ #getServerFaviconFallback(serverUrl: string): string | null { try { @@ -505,27 +499,6 @@ class MCPStore { return null; } - isAnyServerLoading(): boolean { - return this.getServers().some((s) => { - const state = this.getHealthCheckState(s.id); - - return ( - state.status === HealthCheckStatus.IDLE || state.status === HealthCheckStatus.CONNECTING - ); - }); - } - - getServersSorted(): MCPServerSettingsEntry[] { - const servers = this.getServers(); - if (this.isAnyServerLoading()) { - return servers; - } - - return [...servers].sort((a, b) => - this.getServerLabel(a).localeCompare(this.getServerLabel(b)) - ); - } - addServer( serverData: Omit & { id?: string } ): MCPServerSettingsEntry { @@ -579,10 +552,11 @@ class MCPStore { } /** - * MCP servers selectable in chat-add UIs and the settings page. + * MCP servers selectable in chat-add UIs and the settings page, + * in the order they were added to the config. */ get visibleMcpServers(): MCPServerSettingsEntry[] { - return this.getServersSorted().filter((server) => server.enabled); + return this.getServers().filter((server) => server.enabled); } async ensureInitialized(perChatOverrides?: McpServerOverride[]): Promise { diff --git a/tools/ui/src/lib/types/index.ts b/tools/ui/src/lib/types/index.ts index 408ac0cbd..91ceb72e7 100644 --- a/tools/ui/src/lib/types/index.ts +++ b/tools/ui/src/lib/types/index.ts @@ -128,7 +128,6 @@ export type { MCPClientConfig, MCPServerSettingsEntry, MCPServerDisplayInfo, - RecommendedMCPServer, MCPToolCall, OpenAIToolDefinition, ServerStatus, diff --git a/tools/ui/src/lib/types/mcp.d.ts b/tools/ui/src/lib/types/mcp.d.ts index bf7fd3285..dca57d396 100644 --- a/tools/ui/src/lib/types/mcp.d.ts +++ b/tools/ui/src/lib/types/mcp.d.ts @@ -226,15 +226,6 @@ export type MCPServerSettingsEntry = MCPServerDisplayInfo & { useProxy?: boolean; }; -/** - * Pre-defined recommended MCP server shown to the user in onboarding/picker UIs. - */ -export interface RecommendedMCPServer extends MCPServerDisplayInfo { - description: string; - enabled: boolean; - requestTimeoutSeconds: number; -} - export interface MCPHostManagerConfig { servers: MCPClientConfig['servers']; clientInfo?: Implementation; diff --git a/tools/ui/src/routes/+layout.svelte b/tools/ui/src/routes/+layout.svelte index 29b5b1510..f3c024bb9 100644 --- a/tools/ui/src/routes/+layout.svelte +++ b/tools/ui/src/routes/+layout.svelte @@ -8,7 +8,6 @@ import { onMount } from 'svelte'; import { SidebarNavigation, DialogConversationTitleUpdate } from '$lib/components/app'; - import { DialogMcpServerRecommendations } from '$lib/components/app/dialogs'; import { PwaMetaTags, PwaRefreshAlert } from '$lib/components/pwa'; import { pwaAssetsHead } from 'virtual:pwa-assets/head'; @@ -27,7 +26,6 @@ import { FAVICON_PATHS, FAVICON_SELECTORS } from '$lib/constants/pwa'; import { useKeyboardShortcuts } from '$lib/hooks/use-keyboard-shortcuts.svelte'; import { usePwa } from '$lib/hooks/use-pwa.svelte'; - import { useMcpRecommendations } from '$lib/hooks/use-mcp-recommendations.svelte'; import { conversations } from '$lib/stores/conversations.svelte'; import { isMobile } from '$lib/stores/viewport.svelte'; import { theme } from '$lib/stores/theme.svelte'; @@ -39,8 +37,6 @@ let innerHeight = $state(); let innerWidth = $state(browser ? window.innerWidth : 0); - const mcpRecommendations = useMcpRecommendations(); - let chatSidebar: | { activateSearchMode?: () => void; @@ -239,7 +235,10 @@ }); // Background MCP server health checks on app load - // Fetch enabled servers from settings and run health checks in background + // Fetch enabled servers from settings and run health checks in background. + // Only IDLE servers are checked; already-resolved (SUCCESS / ERROR) servers + // keep their existing state, so adding or removing a server does not flash + // every other card back through skeleton state. $effect(() => { if (!browser) return; @@ -251,7 +250,7 @@ if (enabledServers.length > 0) { untrack(() => { // Run health checks in background (don't await) - mcpStore.runHealthChecksForServers(enabledServers, false).catch((error) => { + mcpStore.runHealthChecksForServers(enabledServers, true).catch((error) => { console.warn('[layout] MCP health checks failed:', error); }); }); @@ -325,11 +324,6 @@ onConfirm={handleTitleUpdateConfirm} onCancel={handleTitleUpdateCancel} /> - - diff --git a/tools/ui/tests/unit/mcp-default-overrides-merge.test.ts b/tools/ui/tests/unit/mcp-default-overrides-merge.test.ts new file mode 100644 index 000000000..ce721ca6c --- /dev/null +++ b/tools/ui/tests/unit/mcp-default-overrides-merge.test.ts @@ -0,0 +1,158 @@ +import { afterEach, beforeAll, beforeEach, describe, expect, it } from 'vitest'; +import { STORAGE_APP_NAME, CONFIG_LOCALSTORAGE_KEY } from '$lib/constants'; + +// node env unit project has no DOM, install a minimal localStorage backed by a Map +beforeAll(() => { + const store = new Map(); + const polyfill: Storage = { + get length() { + return store.size; + }, + clear: () => store.clear(), + getItem: (k) => (store.has(k) ? store.get(k)! : null), + key: (i) => Array.from(store.keys())[i] ?? null, + removeItem: (k) => { + store.delete(k); + }, + setItem: (k, v) => { + store.set(k, String(v)); + } + }; + (globalThis as unknown as { localStorage: Storage }).localStorage = polyfill; +}); + +/** + * Migration `mcp-default-overrides-merge-v1` folds the values of the parallel + * `mcpDefaultServerOverrides` config entry onto `mcpServers[i].enabled` (the + * single source of truth for new-chat defaults). The legacy key is kept on + * disk for downgrade compatibility. + */ +describe('mcp-default-overrides-merge-v1 migration', () => { + const MIGRATION_STATE_KEY = `${STORAGE_APP_NAME}.migration-state`; + const MCP_DEFAULT_OVERRIDES_KEY = `${STORAGE_APP_NAME}.mcpDefaultServerOverrides`; + + beforeEach(async () => { + localStorage.clear(); + // Reset the migration run counter so `runAllMigrations` is guaranteed to execute. + await import('$lib/services/migration.service').then((mod) => + mod.MigrationService.resetState() + ); + }); + + afterEach(() => { + localStorage.clear(); + }); + + async function runMigrations() { + const { MigrationService } = await import('$lib/services/migration.service'); + await MigrationService.runAllMigrations(); + } + + function readConfig(): Record { + const raw = localStorage.getItem(CONFIG_LOCALSTORAGE_KEY); + + return raw ? (JSON.parse(raw) as Record) : {}; + } + + function writeConfig(config: Record) { + localStorage.setItem(CONFIG_LOCALSTORAGE_KEY, JSON.stringify(config)); + } + + it('applies matching overrides onto mcpServers[i].enabled and preserves the legacy key', async () => { + writeConfig({ + mcpServers: JSON.stringify([ + { id: 'exa', enabled: false, url: 'https://mcp.exa.ai/mcp' }, + { id: 'hf', enabled: false, url: 'https://huggingface.co/mcp' } + ]), + [MCP_DEFAULT_OVERRIDES_KEY]: JSON.stringify([ + { serverId: 'exa', enabled: true }, + { serverId: 'hf', enabled: false } + ]) + }); + + await runMigrations(); + + const after = readConfig(); + const servers = JSON.parse(after.mcpServers as string) as Array<{ + id: string; + enabled: boolean; + }>; + + expect(servers.find((s) => s.id === 'exa')?.enabled).toBe(true); + expect(servers.find((s) => s.id === 'hf')?.enabled).toBe(false); + expect(MCP_DEFAULT_OVERRIDES_KEY in after).toBe(true); + }); + + it('skips override ids that do not match any configured server', async () => { + writeConfig({ + mcpServers: JSON.stringify([{ id: 'exa', enabled: false, url: 'https://mcp.exa.ai/mcp' }]), + [MCP_DEFAULT_OVERRIDES_KEY]: JSON.stringify([ + { serverId: 'orphan', enabled: true }, + { serverId: 'exa', enabled: true } + ]) + }); + + await runMigrations(); + + const after = readConfig(); + const servers = JSON.parse(after.mcpServers as string) as Array<{ + id: string; + enabled: boolean; + }>; + + expect(servers).toHaveLength(1); + expect(servers[0].enabled).toBe(true); + expect(MCP_DEFAULT_OVERRIDES_KEY in after).toBe(true); + }); + + it('is a no-op when there are no legacy overrides', async () => { + writeConfig({ + mcpServers: JSON.stringify([{ id: 'exa', enabled: true, url: 'https://mcp.exa.ai/mcp' }]) + }); + + await runMigrations(); + + const after = readConfig(); + const servers = JSON.parse(after.mcpServers as string) as Array<{ + id: string; + enabled: boolean; + }>; + + expect(servers[0].enabled).toBe(true); + expect(MCP_DEFAULT_OVERRIDES_KEY in after).toBe(false); + }); + + it('does not rewrite mcpServers when override.enabled already matches', async () => { + const originalServers = JSON.stringify([ + { id: 'exa', enabled: true, url: 'https://mcp.exa.ai/mcp' } + ]); + + writeConfig({ + mcpServers: originalServers, + [MCP_DEFAULT_OVERRIDES_KEY]: JSON.stringify([{ serverId: 'exa', enabled: true }]) + }); + + await runMigrations(); + + const after = readConfig(); + expect(after.mcpServers).toBe(originalServers); + expect(MCP_DEFAULT_OVERRIDES_KEY in after).toBe(true); + }); + + it('records itself as completed so subsequent loads do not re-run', async () => { + writeConfig({ + mcpServers: JSON.stringify([{ id: 'exa', enabled: false, url: 'https://mcp.exa.ai/mcp' }]), + [MCP_DEFAULT_OVERRIDES_KEY]: JSON.stringify([{ serverId: 'exa', enabled: true }]) + }); + + const { MigrationService } = await import('$lib/services/migration.service'); + + await MigrationService.runAllMigrations(); + + const stateRaw = localStorage.getItem(MIGRATION_STATE_KEY); + expect(stateRaw).not.toBeNull(); + const state = JSON.parse(stateRaw!) as { completed: string[]; failed: string[] }; + expect(state.completed).toContain('mcp-default-overrides-merge-v1'); + expect(state.failed).not.toContain('mcp-default-overrides-merge-v1'); + }); +}); diff --git a/tools/ui/tests/unit/mcp-servers-default.test.ts b/tools/ui/tests/unit/mcp-servers-default.test.ts new file mode 100644 index 000000000..35eec6fb7 --- /dev/null +++ b/tools/ui/tests/unit/mcp-servers-default.test.ts @@ -0,0 +1,19 @@ +import { describe, expect, it } from 'vitest'; +import { SETTINGS_KEYS } from '$lib/constants/settings-keys'; + +/** + * Default-value policy for the `MCP_SERVERS` setting. + * + * Earlier versions of the UI preloaded a hard-coded list of suggested + * MCP servers into this setting on first install. That caused silent + * third-party HTTP requests at app load (see issue #25509) and a popup + * "recommendation" dialog (see issue #25274). New users must now opt + * in explicitly when adding a server, so the default is an empty list. + */ +describe('MCP_SERVERS default value', () => { + it('does not preload any servers in the MCP_SERVERS setting default', async () => { + const { SETTING_CONFIG_DEFAULT } = await import('$lib/constants/settings-registry'); + + expect(SETTING_CONFIG_DEFAULT[SETTINGS_KEYS.MCP_SERVERS]).toBe('[]'); + }, 15000); +}); diff --git a/tools/ui/tests/unit/parse-mcp-server-settings.test.ts b/tools/ui/tests/unit/parse-mcp-server-settings.test.ts index 956c677d5..f5e0b3a96 100644 --- a/tools/ui/tests/unit/parse-mcp-server-settings.test.ts +++ b/tools/ui/tests/unit/parse-mcp-server-settings.test.ts @@ -5,11 +5,10 @@ import { DEFAULT_MCP_CONFIG, MCP_SERVER_ID_PREFIX } from '$lib/constants/mcp'; /** * Tests for the mcpServers settings parser. * - * The branch seeds the MCP servers setting with a default value of - * `JSON.stringify(RECOMMENDED_MCP_SERVERS)`, so the parser has to be - * resilient to anything that may live in the user's localStorage: malformed - * JSON, wrong shapes, missing fields, falsy-but-not-zero numbers, and entry - * arrays that have been mutated by the user via the settings form. + * The parser has to be resilient to anything that may live in the + * user's localStorage: malformed JSON, wrong shapes, missing fields, + * falsy-but-not-zero numbers, and entry arrays that have been mutated + * by the user via the settings form. */ describe('parseMcpServerSettings', () => { it('returns an empty array for falsy or whitespace-only input', () => { diff --git a/tools/ui/tests/unit/recommended-mcp-servers.test.ts b/tools/ui/tests/unit/recommended-mcp-servers.test.ts deleted file mode 100644 index 3f6fd8f11..000000000 --- a/tools/ui/tests/unit/recommended-mcp-servers.test.ts +++ /dev/null @@ -1,90 +0,0 @@ -import { describe, expect, it } from 'vitest'; -import { - RECOMMENDED_MCP_SERVER_IDS, - RECOMMENDED_MCP_SERVERS -} from '$lib/constants/recommended-mcp-servers'; -import { parseMcpServerSettings } from '$lib/utils/mcp'; -import { DEFAULT_MCP_CONFIG, MCP_SERVER_ID_PREFIX } from '$lib/constants/mcp'; - -/** - * Tests for the predefined recommended MCP servers. - * - * These are surfaced to first-time users via - * DialogMcpServerRecommendations and used as the default value of the MCP - * servers setting, so a regression that breaks the round-trip through the - * settings parser would silently break onboarding for new users. - */ -describe('RECOMMENDED_MCP_SERVERS', () => { - it('lists at least one entry and uses stable, unique ids', () => { - expect(RECOMMENDED_MCP_SERVERS.length).toBeGreaterThan(0); - - const ids = RECOMMENDED_MCP_SERVERS.map((server) => server.id); - expect(new Set(ids).size).toBe(ids.length); - - for (const id of ids) { - expect(id).toMatch(/^[a-z0-9-]+$/); - expect(id.toLowerCase()).not.toContain(MCP_SERVER_ID_PREFIX.toLowerCase()); - } - }); - - it('requires a name, description and url for every entry', () => { - for (const server of RECOMMENDED_MCP_SERVERS) { - expect(server.name?.trim().length ?? 0).toBeGreaterThan(0); - expect(server.description.trim().length).toBeGreaterThan(0); - expect(server.url.trim().length).toBeGreaterThan(0); - expect(() => new URL(server.url)).not.toThrow(); - } - }); -}); - -describe('RECOMMENDED_MCP_SERVER_IDS', () => { - it('matches the ids declared in RECOMMENDED_MCP_SERVERS', () => { - expect(RECOMMENDED_MCP_SERVER_IDS.size).toBe(RECOMMENDED_MCP_SERVERS.length); - - for (const server of RECOMMENDED_MCP_SERVERS) { - expect(RECOMMENDED_MCP_SERVER_IDS.has(server.id)).toBe(true); - } - }); -}); - -describe('recommended-mcp-servers default value', () => { - it('round-trips cleanly through parseMcpServerSettings', () => { - const serialized = JSON.stringify(RECOMMENDED_MCP_SERVERS); - const parsed = parseMcpServerSettings(serialized); - - expect(parsed).toHaveLength(RECOMMENDED_MCP_SERVERS.length); - - for (let index = 0; index < RECOMMENDED_MCP_SERVERS.length; index++) { - const source = RECOMMENDED_MCP_SERVERS[index]; - const entry = parsed[index]; - - expect(entry).toBeDefined(); - expect(entry?.id).toBe(source.id); - expect(entry?.url).toBe(source.url); - expect(entry?.enabled).toBe(source.enabled); - expect(entry?.requestTimeoutSeconds).toBe(source.requestTimeoutSeconds); - expect(entry?.name).toBe(source.name); - - // Headers and useProxy are not set on recommended servers; the - // parser must fall back to the inactive defaults rather than - // surfacing undefined-boundary states. - expect(entry?.headers).toBeUndefined(); - expect(entry?.useProxy).toBe(false); - } - }); - - it('uses the global default timeout when one is not specified on an entry', () => { - const sourceOnlyRequired = { - id: 'roundtrip-only', - name: 'Only required fields', - url: 'https://example.test/mcp', - description: 'Smoke entry for parser roundtrip with default timeout.', - enabled: true - }; - - const parsed = parseMcpServerSettings(JSON.stringify([sourceOnlyRequired])); - const entry = parsed[0]; - - expect(entry?.requestTimeoutSeconds).toBe(DEFAULT_MCP_CONFIG.requestTimeoutSeconds); - }); -}); diff --git a/vendor/cpp-httplib/httplib.cpp b/vendor/cpp-httplib/httplib.cpp index d65b7921b..e58fb3eea 100644 --- a/vendor/cpp-httplib/httplib.cpp +++ b/vendor/cpp-httplib/httplib.cpp @@ -3705,6 +3705,12 @@ write_content_chunked(Stream &strm, const ContentProvider &content_provider, // Trailer if (trailer) { 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; + } std::string field_line = kv.first + ": " + kv.second + "\r\n"; if (!write_data(strm, field_line.data(), field_line.size())) { ok = false; @@ -8301,8 +8307,8 @@ void Server::apply_ranges(const Request &req, Response &res, } } - auto length = std::to_string(res.body.size()); - res.set_header("Content-Length", length); + res.content_length_ = res.body.size(); + res.set_header("Content-Length", std::to_string(res.content_length_)); } } @@ -10270,6 +10276,11 @@ Result ClientImpl::Get(const std::string &path, return Get(path, Headers(), std::move(progress)); } +Result ClientImpl::Get(const std::string &path, const Params ¶ms, + DownloadProgress progress) { + return Get(path, params, Headers(), std::move(progress)); +} + Result ClientImpl::Get(const std::string &path, const Params ¶ms, const Headers &headers, DownloadProgress progress) { @@ -11348,6 +11359,10 @@ Result Client::Get(const std::string &path, const Headers &headers, return cli_->Get(path, headers, std::move(response_handler), std::move(content_receiver), std::move(progress)); } +Result Client::Get(const std::string &path, const Params ¶ms, + DownloadProgress progress) { + return cli_->Get(path, params, std::move(progress)); +} Result Client::Get(const std::string &path, const Params ¶ms, const Headers &headers, DownloadProgress progress) { return cli_->Get(path, params, headers, std::move(progress)); @@ -12076,11 +12091,18 @@ bool SSLServer::update_certs_pem(const char *cert_pem, // SSL HTTP client implementation SSLClient::~SSLClient() { - if (ctx_) { tls::free_context(ctx_); } // Make sure to shut down SSL since shutdown_ssl will resolve to the // base function rather than the derived function once we get to the // base class destructor, and won't free the SSL (causing a leak). + // This must happen before the context is freed below: some backends + // (e.g. mbedTLS) have the SSL session borrow a raw pointer into the + // context, so freeing the context first leaves close_notify reading + // freed memory. shutdown_ssl_impl(socket_, true); + if (ctx_) { + tls::free_context(ctx_); + ctx_ = nullptr; + } } bool SSLClient::is_valid() const { return ctx_ != nullptr; } @@ -16501,6 +16523,11 @@ WebSocketClient::~WebSocketClient() { bool WebSocketClient::is_valid() const { return is_valid_; } void WebSocketClient::shutdown_and_close() { + // Send the close frame while the TLS session is still alive: ws_ holds an + // SSLSocketStream that keeps a raw pointer to tls_session_, so the session + // must outlive ws_->close() and ws_.reset() to avoid a use-after-free. + if (ws_ && ws_->is_open()) { ws_->close(); } + ws_.reset(); #ifdef CPPHTTPLIB_SSL_ENABLED if (is_ssl_) { if (tls_session_) { @@ -16510,8 +16537,6 @@ void WebSocketClient::shutdown_and_close() { } } #endif - if (ws_ && ws_->is_open()) { ws_->close(); } - ws_.reset(); if (sock_ != INVALID_SOCKET) { detail::shutdown_socket(sock_); detail::close_socket(sock_); diff --git a/vendor/cpp-httplib/httplib.h b/vendor/cpp-httplib/httplib.h index e7ef56370..de026a253 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.49.0" -#define CPPHTTPLIB_VERSION_NUM "0x003100" +#define CPPHTTPLIB_VERSION "0.50.1" +#define CPPHTTPLIB_VERSION_NUM "0x003201" #ifdef _WIN32 #if defined(_WIN32_WINNT) && _WIN32_WINNT < 0x0A00 @@ -2219,6 +2219,7 @@ public: Result Get(const std::string &path, const Headers &headers, DownloadProgress progress = nullptr); Result Get(const std::string &path, const Headers &headers, ContentReceiver content_receiver, DownloadProgress progress = nullptr); Result Get(const std::string &path, const Headers &headers, ResponseHandler response_handler, ContentReceiver content_receiver, DownloadProgress progress = nullptr); + Result Get(const std::string &path, const Params ¶ms, DownloadProgress progress = nullptr); Result Get(const std::string &path, const Params ¶ms, const Headers &headers, DownloadProgress progress = nullptr); Result Get(const std::string &path, const Params ¶ms, const Headers &headers, ContentReceiver content_receiver, DownloadProgress progress = nullptr); Result Get(const std::string &path, const Params ¶ms, const Headers &headers, ResponseHandler response_handler, ContentReceiver content_receiver, DownloadProgress progress = nullptr); @@ -2602,6 +2603,7 @@ public: Result Get(const std::string &path, const Headers &headers, DownloadProgress progress = nullptr); Result Get(const std::string &path, const Headers &headers, ContentReceiver content_receiver, DownloadProgress progress = nullptr); Result Get(const std::string &path, const Headers &headers, ResponseHandler response_handler, ContentReceiver content_receiver, DownloadProgress progress = nullptr); + Result Get(const std::string &path, const Params ¶ms, DownloadProgress progress = nullptr); Result Get(const std::string &path, const Params ¶ms, const Headers &headers, DownloadProgress progress = nullptr); Result Get(const std::string &path, const Params ¶ms, const Headers &headers, ContentReceiver content_receiver, DownloadProgress progress = nullptr); Result Get(const std::string &path, const Params ¶ms, const Headers &headers, ResponseHandler response_handler, ContentReceiver content_receiver, DownloadProgress progress = nullptr);