From 86d996eb3afaa602c13b1cec1a4960569a6b0c65 Mon Sep 17 00:00:00 2001 From: alex Date: Wed, 6 May 2026 22:07:39 +0200 Subject: [PATCH] fix: ensure prompt cache is invalidated whenever KV cache is cleared (#486) The prompt prefix cache in /v1/chat/completions was broken in several ways: 1. Non-streaming path unconditionally cleared context and reset cache, making prefix caching impossible for non-streaming requests. 2. update_checksum() included tool messages in its checksum calculation while can_use_cache() skipped them, causing permanent cache misses whenever tools were present in the conversation. 3. Error paths (max length, exceptions) in handle_openai_chat_completion cleared the KV cache via clear_context() but did not call prompt_cache.reset(), so the next request would stale cache-hit against an empty KV cache and only insert the last message. 4. Other handlers (/api/chat, /api/generate, /v1/completions) all call clear_context() on every request but never invalidated the prompt cache, causing the same stale-hit problem for any interleaved requests. Fix all cases by introducing a reset_context() helper on RestHandler that couples clear_context() with prompt_cache.reset(), ensuring the prompt cache can never outlive the KV state it refers to. The only bare clear_context() calls that remain are the two cache-miss paths in handle_openai_chat_completion where the cache is immediately re-initialized with update_checksum(). --- src/include/prompt_cache.hpp | 13 +++---- src/server/rest_handler.cpp | 68 +++++++++++++++++++----------------- src/server/rest_handler.hpp | 1 + 3 files changed, 43 insertions(+), 39 deletions(-) diff --git a/src/include/prompt_cache.hpp b/src/include/prompt_cache.hpp index 1760673d..544694d1 100644 --- a/src/include/prompt_cache.hpp +++ b/src/include/prompt_cache.hpp @@ -109,11 +109,11 @@ class PromptCache { bool has_tool_call = msg.contains("tool_calls"); bool skip_this_message = (role == "tool") || has_tool_call; if (template_type == chat_template_type_t::harmony) { - skip_this_message = false; + skip_this_message = false; } if (skip_this_message) continue; - + if(i < messages.size() - 2) check_sum_to_compare = _calculate_checksum(content.data(), content.size(), check_sum_to_compare); new_checksum = _calculate_checksum(content.data(), content.size(), new_checksum); @@ -138,16 +138,17 @@ class PromptCache { uint64_t new_checksum = 0; for (size_t i = 0; i < messages.size(); ++i) { const auto& msg = messages[i]; + const std::string role = msg.value("role", ""); const std::string content = msg.value("content", ""); + bool has_tool_call = msg.contains("tool_calls"); + bool skip_this_message = (role == "tool") || has_tool_call; + if (skip_this_message) continue; new_checksum = _calculate_checksum(content.data(), content.size(), new_checksum); } checksum_ = new_checksum; } - /// @brief Reset the checksum to force cache miss - /// @note This function increments the checksum value by 1 to ensure that - /// the next call to can_use_cache will result in a cache miss. void reset() { checksum_ = checksum_ + 1; } -}; \ No newline at end of file +}; diff --git a/src/server/rest_handler.cpp b/src/server/rest_handler.cpp index 9bce0c5a..4a2006ac 100644 --- a/src/server/rest_handler.cpp +++ b/src/server/rest_handler.cpp @@ -308,6 +308,12 @@ void RestHandler::ensure_embed_model_loaded(const std::string& model_tag) { #endif } +///@brief Clear KV cache and invalidate prompt cache +void RestHandler::reset_context() { + auto_chat_engine->clear_context(); + prompt_cache.reset(); +} + ///@brief Configure chat engine parameters from options and request ///@param options the options JSON object ///@param request the request JSON object @@ -475,13 +481,13 @@ void RestHandler::handle_generate(const json& request, if (!success){ json error_response = {{"error", "Max length reached"}}; send_response(error_response); - this->auto_chat_engine->clear_context(); + this->reset_context(); return; } } catch (const std::exception& e) { json error_response = {{"error", e.what()}}; send_response(error_response); - this->auto_chat_engine->clear_context(); + this->reset_context(); return; } try { @@ -489,7 +495,7 @@ void RestHandler::handle_generate(const json& request, } catch (const std::exception& e) { json error_response = {{"error", e.what()}}; send_response(error_response); - this->auto_chat_engine->clear_context(); + this->reset_context(); return; } auto total_end_time = time_utils::now(); @@ -508,13 +514,13 @@ void RestHandler::handle_generate(const json& request, if (!success){ json error_response = {{"error", "Max length reached"}}; send_response(error_response); - this->auto_chat_engine->clear_context(); + this->reset_context(); return; } } catch (const std::exception& e) { json error_response = {{"error", e.what()}}; send_response(error_response); - this->auto_chat_engine->clear_context(); + this->reset_context(); return; } try { @@ -522,7 +528,7 @@ void RestHandler::handle_generate(const json& request, } catch (const std::exception& e) { json error_response = {{"error", e.what()}}; send_response(error_response); - this->auto_chat_engine->clear_context(); + this->reset_context(); return; } std::string response_text = ss.str(); @@ -591,13 +597,13 @@ void RestHandler::handle_chat(const json& request, if (!success){ json error_response = {{"error", "Max length reached"}}; send_response(error_response); - this->auto_chat_engine->clear_context(); + this->reset_context(); return; } } catch (const std::exception& e) { json error_response = {{"error", e.what()}}; send_response(error_response); - this->auto_chat_engine->clear_context(); + this->reset_context(); return; } try { @@ -605,22 +611,22 @@ void RestHandler::handle_chat(const json& request, if (!success){ json error_response = {{"error", "Max length reached"}}; send_response(error_response); - this->auto_chat_engine->clear_context(); + this->reset_context(); return; } } catch (const std::exception& e) { json error_response = {{"error", e.what()}}; send_response(error_response); - this->auto_chat_engine->clear_context(); + this->reset_context(); return; } auto total_end_time = time_utils::now(); meta_info.total_duration = (uint64_t)time_utils::duration_ns(total_start_time, total_end_time).first; - + ostream.finalize_chat(meta_info); // auto history = this->chat_engine->get_history(); // std::cout << "history: " << history.first << std::endl; - this->auto_chat_engine->clear_context(); + this->reset_context(); } else { // Non-streaming response uniformed_input.messages = messages; @@ -633,7 +639,7 @@ void RestHandler::handle_chat(const json& request, } catch (const std::exception& e) { json error_response = {{"error", e.what()}}; send_response(error_response); - this->auto_chat_engine->clear_context(); + this->reset_context(); return; } //std::string response_text = chat_engine->generate_with_prompt(meta_info, prompts, length_limit, std::cout, payload); @@ -660,7 +666,7 @@ void RestHandler::handle_chat(const json& request, // auto history = this->chat_engine->get_history(); // std::cout << "history: " << history.first << std::endl; - this->auto_chat_engine->clear_context(); + this->reset_context(); } } catch (const std::exception& e) { json error_response = {{"error", e.what()}}; @@ -974,8 +980,7 @@ void RestHandler::handle_openai_chat_completion(const json& request, meta_info.stop_reason = CANCEL_DETECTED; header_print("❌ ", "Prefill Cancelled!"); ostream.finalize(meta_info); - this->auto_chat_engine->clear_context(); - this->prompt_cache.reset(); + this->reset_context(); return; } @@ -987,13 +992,13 @@ void RestHandler::handle_openai_chat_completion(const json& request, }} }; send_response(error_response); - this->auto_chat_engine->clear_context(); + this->reset_context(); return; } } catch (const std::exception& e) { json error_response = {{"error", e.what()}}; send_response(error_response); - this->auto_chat_engine->clear_context(); + this->reset_context(); return; } header_print("FLM", "Start generating..."); @@ -1002,7 +1007,7 @@ void RestHandler::handle_openai_chat_completion(const json& request, } catch (const std::exception& e) { json error_response = {{"error", e.what()}}; send_response(error_response); - this->auto_chat_engine->clear_context(); + this->reset_context(); return; } if (meta_info.stop_reason == CANCEL_DETECTED) { @@ -1013,7 +1018,6 @@ void RestHandler::handle_openai_chat_completion(const json& request, ostream.finalize(meta_info); } else { - this->auto_chat_engine->clear_context(); nullstream nstream; json response; std::string response_text; @@ -1025,8 +1029,7 @@ void RestHandler::handle_openai_chat_completion(const json& request, meta_info.stop_reason = CANCEL_DETECTED; header_print("❌ ", "Prefill Cancelled!"); send_response(response); - this->auto_chat_engine->clear_context(); - this->prompt_cache.reset(); + this->reset_context(); return; } json error_response = { @@ -1037,13 +1040,13 @@ void RestHandler::handle_openai_chat_completion(const json& request, }} }; send_response(error_response); - this->auto_chat_engine->clear_context(); + this->reset_context(); return; } } catch (const std::exception& e) { json error_response = {{"error", e.what()}}; send_response(error_response); - this->auto_chat_engine->clear_context(); + this->reset_context(); return; } header_print("FLM", "Start generating..."); @@ -1052,7 +1055,7 @@ void RestHandler::handle_openai_chat_completion(const json& request, } catch (const std::exception& e) { json error_response = {{"error", e.what()}}; send_response(error_response); - this->auto_chat_engine->clear_context(); + this->reset_context(); return; } // check response_text @@ -1080,7 +1083,6 @@ void RestHandler::handle_openai_chat_completion(const json& request, header_print("❌ ", "Generation Cancelled!"); } send_response(response); - this->prompt_cache.reset(); } } catch (const std::exception& e) { @@ -1207,13 +1209,13 @@ void RestHandler::handle_openai_completion(const json& request, if (!success) { json error_response = { {"error", "Max length reached"} }; send_response(error_response); - this->auto_chat_engine->clear_context(); + this->reset_context(); return; } } catch (const std::exception& e) { json error_response = {{"error", e.what()}}; send_response(error_response); - this->auto_chat_engine->clear_context(); + this->reset_context(); return; } try { @@ -1221,12 +1223,12 @@ void RestHandler::handle_openai_completion(const json& request, } catch (const std::exception& e) { json error_response = {{"error", e.what()}}; send_response(error_response); - this->auto_chat_engine->clear_context(); + this->reset_context(); return; } ostream.finalize(meta_info); - this->auto_chat_engine->clear_context(); + this->reset_context(); } else { std::stringstream ss; @@ -1238,13 +1240,13 @@ void RestHandler::handle_openai_completion(const json& request, if (!success) { json error_response = { {"error", "Max length reached"} }; send_response(error_response); - this->auto_chat_engine->clear_context(); + this->reset_context(); return; } } catch (const std::exception& e) { json error_response = {{"error", e.what()}}; send_response(error_response); - this->auto_chat_engine->clear_context(); + this->reset_context(); return; } try { @@ -1252,7 +1254,7 @@ void RestHandler::handle_openai_completion(const json& request, } catch (const std::exception& e) { json error_response = {{"error", e.what()}}; send_response(error_response); - this->auto_chat_engine->clear_context(); + this->reset_context(); return; } std::string response_text = ss.str(); diff --git a/src/server/rest_handler.hpp b/src/server/rest_handler.hpp index 98b1ab82..a180bba6 100644 --- a/src/server/rest_handler.hpp +++ b/src/server/rest_handler.hpp @@ -111,6 +111,7 @@ class RestHandler { bool ensure_model_loaded(const std::string& model_tag); void ensure_asr_model_loaded(const std::string& model_tag); void ensure_embed_model_loaded(const std::string& model_tag); + void reset_context(); void configure_chat_engine_parameters(const json& options, const json& request); json build_nstream_response(std::string response_text);