diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp
index fbec0049..eea5165f 100644
--- a/tools/server/server-task.cpp
+++ b/tools/server/server-task.cpp
@@ -2503,44 +2503,37 @@ bool server_prompt_cache::load_entry(std::list<server_prompt_cache_state>::itera
const bool learned = is_learned(it->prompt.meta);
const bool consume = !keep_entry && !learned;
- {
- auto & data = it->data.main;
-
- const size_t size = data.size();
- const size_t n = llama_state_seq_set_data_ext(ctx_tgt, data.data(), size, id_slot, 0);
- if (n != size) {
- SRV_ERR("failed to restore state with size %zu\n", size);
+ // an empty state (a learned checkpoint without its state) restores nothing; llama reports 0 bytes read for it,
+ // which would otherwise pass as a complete restore
+ if (it->data.main.empty()) {
+ SRV_WRN("prompt cache: entry %s (%s, %d tokens) has no state to restore\n", hex_id(it->prompt.meta.id).c_str(),
+ cache_source_name(it->prompt.meta.source), it->prompt.n_tokens());
+ return false;
+ }
+ {
+ const auto & data = it->data.main;
+ const size_t n = llama_state_seq_set_data_ext(ctx_tgt, data.data(), data.size(), id_slot, 0);
+ if (n != data.size()) {
+ SRV_ERR("failed to restore state with size %zu\n", data.size());
return false;
}
-
- if (consume) {
- data.clear();
- data.shrink_to_fit();
- }
}
{
- auto & data = it->data.drft;
-
+ const auto & data = it->data.drft;
if (!data.empty()) {
GGML_ASSERT(ctx_dft);
- const size_t size = data.size();
- const size_t n = llama_state_seq_set_data_ext(ctx_dft, data.data(), size, id_slot, 0);
- if (n != size) {
- SRV_WRN("failed to restore state with size %zu\n", size);
-
+ const size_t n = llama_state_seq_set_data_ext(ctx_dft, data.data(), data.size(), id_slot, 0);
+ if (n != data.size()) {
+ SRV_WRN("failed to restore state with size %zu\n", data.size());
return false;
}
-
- if (consume) {
- data.clear();
- data.shrink_to_fit();
- }
}
}
+ // consumed only once both restores succeeded: a failed draft restore must not leave an entry without its state
if (consume) {
prompt = std::move(it->prompt);
@@ -2725,7 +2718,7 @@ bool server_prompt_cache::family_class(uint64_t family, cache_class & out) const
void server_prompt_cache::apply_family(cache_meta & m) const {
cache_class cls;
- if (m.source != CACHE_SOURCE_EXPLICIT && family_class(m.family, cls)) {
+ if (m.source != CACHE_SOURCE_EXPLICIT && !is_learned(m) && family_class(m.family, cls)) {
m.cls = cls;
m.source = CACHE_SOURCE_FAMILY;
}
@@ -2947,7 +2940,8 @@ bool server_prompt_cache::set_entry(uint64_t id, cache_class cls, int64_t ttl) {
bool found = false;
auto apply = [&](cache_meta & m) {
m.cls = cls;
- m.source = CACHE_SOURCE_EXPLICIT;
+ // a learned prefix stays one: only its class and ttl change
+ m.source = is_learned(m) ? CACHE_SOURCE_LEARNED : CACHE_SOURCE_EXPLICIT;
m.ttl = ttl;
m.last_used = std::max(m.last_used, unix_now());
set_expiry(m);