fix(pflash): drafter resumes after a tool output

server/src/common/model_backend.h

@@ -490,6 +490,9 @@ struct ModelBackend {
         // Earlier user questions (their scorer windows), most recent first:
         // they score the context alongside the query at halving weights.
         std::vector<PFlashTokenSpan> history_query_spans;
+        // Windows a later request scores as history queries, never scored
+        // here: the scorer keeps their rows so that request can resume.
+        std::vector<PFlashTokenSpan> keep_query_spans;
         // The latest user turn's tail, a second query window at full weight
         // (prompt-end chat queries; {-1, -1} otherwise).
         PFlashTokenSpan turn_query_span{-1, -1};

server/src/pflash/pflash_drafter.cpp

@@ -114,7 +114,8 @@ std::vector<int32_t> drafter_score_and_compress(
     const std::vector<PFlashTokenSpan> & history_queries,
     PFlashTokenSpan turn_query,
     const std::vector<PFlashTokenSpan> & extra_queries,
-    const PFlashRequestMessages * request_messages) {
+    const PFlashRequestMessages * request_messages,
+    const std::vector<PFlashTokenSpan> & keep_queries) {
     pflash_clear_kept_spans();
     if (!ctx.loaded) {
         set_last_error("drafter not loaded");
@@ -141,6 +142,12 @@ std::vector<int32_t> drafter_score_and_compress(
                 experiment.history_queries.push_back(window);
             }
         }
+        for (const auto & window : keep_queries) {
+            if (window.begin >= 0 && window.end > window.begin &&
+                window.end <= (int) ids.size()) {
+                experiment.keep_queries.push_back(window);
+            }
+        }
         if (turn_query.begin >= 0 && turn_query.end > turn_query.begin &&
             turn_query.end <= (int) ids.size()) {
             experiment.turn_query = turn_query;

server/src/pflash/pflash_drafter.h

@@ -100,6 +100,9 @@ void free_drafter_weights(DrafterContext & ctx);
 //                    messages' content spans the request probe reads
 //                    (nullptr: none, the request falls back to today's
 //                    windows)
+//   keep_queries strict selection only: windows a later request scores as
+//                    history queries; never scored here, their rows are
+//                    kept so that request can resume the scoring session
 //
 // On failure returns empty vector + sets last_error.
 std::vector<int32_t> drafter_score_and_compress(
@@ -116,6 +119,7 @@ std::vector<int32_t> drafter_score_and_compress(
     const std::vector<PFlashTokenSpan> & history_queries = {},
     PFlashTokenSpan turn_query = {-1, -1},
     const std::vector<PFlashTokenSpan> & extra_queries = {},
-    const PFlashRequestMessages * request_messages = nullptr);
+    const PFlashRequestMessages * request_messages = nullptr,
+    const std::vector<PFlashTokenSpan> & keep_queries = {});
 
 } // namespace luce::common

server/src/pflash/pflash_scoring_resume.cpp

@@ -24,6 +24,12 @@ int pflash_scoring_resume(int checkpoint, int shared, int S, int cap) {
     return checkpoint;
 }
 
+std::vector<int> pflash_span_tokens(int begin, int end) {
+    std::vector<int> tokens;
+    for (int t = begin; t < end; ++t) tokens.push_back(t);
+    return tokens;
+}
+
 bool PFlashExactRowSets::put(const std::vector<int> & tokens, const float * rows, int hidden, int exact_end) {
     if (tokens.empty() || hidden <= 0 || !rows || tokens.front() < 0 || tokens.back() >= exact_end) return false;
     for (size_t i = 1; i < tokens.size(); ++i) {

server/src/pflash/pflash_scoring_resume.h

@@ -45,6 +45,9 @@ int pflash_session_capacity(int S, int ctx_limit);
 // the two that are > 0, or 0 when neither is.
 int pflash_session_ctx_limit(int drafter_ctx, int engine_max_ctx);
 
+// The token positions of the window [begin, end), ascending; empty when end <= begin.
+std::vector<int> pflash_span_tokens(int begin, int end);
+
 // Exact F32 block-14 rows of query token sets (the query step's candidate and END sets), kept per session for later
 // turns whose activations are gone. Only sets whose last token is below the storing call's exact-rows end are put.
 class PFlashExactRowSets {

server/src/pflash/pflash_selection.h

@@ -139,6 +139,12 @@ struct PFlashSelectionConfig {
     // first. The head scores the context against each and mixes the masses
     // with the query's at weights 1/2, 1/4, ... (multi-turn chats).
     std::vector<luce::common::PFlashTokenSpan> history_queries;
+    // Per request: windows a later request of the same chat will score as
+    // history queries but this one does not (the latest user question's
+    // counterpart in a tool session, whose prompt ends after tool output).
+    // Never scored: the scorer keeps their rows in its session, so the next
+    // request can resume instead of scoring from token 0.
+    std::vector<luce::common::PFlashTokenSpan> keep_queries;
     // Per request: the tail of the latest user turn, scored as a second
     // query window at full weight next to the prompt-end query. The last
     // token reads the whole request; the user's own tokens match literal

server/src/pflash/qwen35_drafter.cpp

@@ -684,13 +684,9 @@ Qwen35ScoringSession * acquire_scoring_session(
     // Stored windows hold exact rows only (see the store below): one is
     // reusable while it sits in the shared prefix.
     const auto stored_window = [&](const Qwen35ScoringSession & session, int begin, int end, int shared) {
-        const size_t floats = (size_t) w.n_embd * (size_t) (end - begin);
-        for (const auto & window : session.query_windows) {
-            if (window.begin == begin && window.end == end && end <= shared && window.rows.size() == floats) {
-                return true;
-            }
-        }
-        return false;
+        std::vector<float> rows;
+        return end <= shared &&
+               session.query_windows.lookup(luce::pflash::pflash_span_tokens(begin, end), w.n_embd, rows);
     };
     Qwen35ScoringSession * best = nullptr;
     int best_shared = 0;
@@ -1205,6 +1201,7 @@ static std::vector<int32_t> strict_score_impl(
     const int request_rows_invalidated = request_probe
         ? session->request_rows.invalidate_from(shared_prefix) : 0;
     session->query_sets.invalidate_from(shared_prefix);
+    session->query_windows.invalidate_from(shared_prefix);
     // A session is released (freed or kept) on every exit below.
     struct SessionExit {
         std::unique_ptr<Qwen35ScoringSession> & scratch;
@@ -1411,14 +1408,8 @@ static std::vector<int32_t> strict_score_impl(
                                     rows.size() * sizeof(float));
             return true;
         }
-        for (const auto & stored : session->query_windows) {
-            if (stored.begin == begin && stored.end == end &&
-                end <= shared_prefix && stored.rows.size() == rows.size()) {
-                rows = stored.rows;
-                return true;
-            }
-        }
-        return false;
+        return end <= shared_prefix &&
+               session->query_windows.lookup(luce::pflash::pflash_span_tokens(begin, end), hidden, rows);
     };
     {
         ScoredWindow query;
@@ -1462,6 +1453,16 @@ static std::vector<int32_t> strict_score_impl(
             }
         }
     }
+    // Rows a later request scores as history queries: kept with the session,
+    // never scored here.
+    std::vector<ScoredWindow> keep_rows;
+    for (const auto & span : experiment.keep_queries) {
+        if (drop_history || span.end > query_start || span.end - span.begin < 1) continue;
+        ScoredWindow keep;
+        keep.begin = span.begin;
+        keep.end = span.end;
+        if (rows_for(span.begin, span.end, keep.rows)) keep_rows.push_back(std::move(keep));
+    }
 
     ggml_tensor * wk_src = st.head_loaded ? st.head_wk : L.wk;
     session->keys_trained = st.head_loaded;
@@ -2224,15 +2225,13 @@ static std::vector<int32_t> strict_score_impl(
     for (const auto & [tokens, rows] : qrows) {
         session->query_sets.put(tokens, rows.data(), hidden, exact_rows_end);
     }
-    for (const auto & window : windows) {
-        if (window.end > exact_rows_end) continue;
-        auto & stored = session->query_windows;
-        stored.erase(std::remove_if(stored.begin(), stored.end(),
-            [&window] (const Qwen35ScoringSession::QueryRows & old) {
-                return old.begin == window.begin && old.end == window.end;
-            }), stored.end());
-        stored.push_back({window.begin, window.end, window.rows});
-        if (stored.size() > 8) stored.erase(stored.begin());
+    // The kept windows go last: the next request needs them, so the cap
+    // evicts this call's windows first.
+    for (const auto * list : {&windows, &keep_rows}) {
+        for (const auto & window : *list) {
+            session->query_windows.put(luce::pflash::pflash_span_tokens(window.begin, window.end),
+                                       window.rows.data(), hidden, exact_rows_end);
+        }
     }
     session->last_used = ++st.session_clock;
 

server/src/pflash/qwen35_drafter.h

@@ -63,12 +63,8 @@ struct Qwen35ScoringSession {
     // questions), reused while they sit in the shared prefix: an agent step
     // appends tool output after the same user turn, and a new turn's history
     // queries are earlier turns' queries. Most recent last, at most 8.
-    struct QueryRows {
-        int begin = -1;
-        int end = -1;
-        std::vector<float> rows;            // [hidden, end - begin]
-    };
-    std::vector<QueryRows> query_windows;
+    // Keyed by the window's token positions.
+    luce::pflash::PFlashExactRowSets query_windows;
     // The query step's (PFLASH_QUERY_CANDIDATES / PFLASH_QUERY_PIN) query token sets, F32 rows exact for later
     // prompts that share them (never the BF16 request_rows).
     luce::pflash::PFlashExactRowSets query_sets;

server/src/qwen35/qwen35_backend.cpp

@@ -1513,7 +1513,7 @@ std::vector<ModelBackend::CompressResult> Qwen35Backend::compress_batch(
             score_query_end, request.required_instruction_spans,
             request.query_suffix_candidates, request.history_query_spans,
             request.turn_query_span, request.extra_query_spans,
-            &request.request_messages);
+            &request.request_messages, request.keep_query_spans);
         result.ok = !result.compressed_ids.empty();
         if (result.ok) result.kept_spans = pflash_last_kept_spans();
         if (result.ok) {

server/src/qwen35/qwen35_layer_split_adapter.cpp

@@ -1396,7 +1396,8 @@ Qwen35LayerSplitAdapter::compress(const ModelBackend::CompressRequest & req) {
         /*chunk_size=*/32, req.score_query_tokens, /*pool_kernel=*/13,
         score_query_end, req.required_instruction_spans,
         req.query_suffix_candidates, req.history_query_spans,
-            req.turn_query_span, req.extra_query_spans, &req.request_messages);
+            req.turn_query_span, req.extra_query_spans, &req.request_messages,
+            req.keep_query_spans);
     result.ok = !result.compressed_ids.empty();
     if (result.ok) result.kept_spans = pflash_last_kept_spans();
     if (result.ok) {

server/src/server/http_server.cpp

@@ -5008,6 +5008,9 @@ std::string HttpServer::apply_pflash_compression(
     // Earlier user questions of a multi-turn chat, most recent first: they
     // score the context alongside the current query at halving weights.
     std::vector<PFlashTokenSpan> history_query_spans;
+    // The window a later request of this chat scores for the latest question
+    // as a history query, when this request does not score it (below).
+    std::vector<PFlashTokenSpan> keep_query_spans;
     // The latest user turn's tail (PFLASH_SELECT_QUERY_TOKENS), a second
     // query window next to the prompt-end query: literal strings in the
     // question (an identifier, a described function) match their passage.
@@ -5399,6 +5402,22 @@ std::string HttpServer::apply_pflash_compression(
                                 {window.end - window.tokens, window.end});
                         }
                     }
+                    // A tool session's prompt ends after tool output, so the
+                    // latest question's history window (the header of the
+                    // reply that followed it) is no query here. The next
+                    // question scores it; the drafter keeps its rows so that
+                    // request resumes the scoring session instead of
+                    // scoring the whole prompt again.
+                    const size_t reply = (size_t) chat_turn.query_turn + 1;
+                    if (prompt_end_query && history_queries > 0 &&
+                        chat_turn.query_turn >= 0 &&
+                        reply < chat_turn.turns.size() &&
+                        chat_turn.turns[reply].role != "user" &&
+                        chat_turn.turns[reply].content_begin >
+                            chat_turn.turns[reply].role_begin) {
+                        const int end = chat_turn.turns[reply].content_begin;
+                        keep_query_spans.push_back({end - 1, end});
+                    }
                 }
                 // The last non-system turns before the generation prompt
                 // (latest first), from the rendered prompt's chat markers:
@@ -5835,6 +5854,7 @@ std::string HttpServer::apply_pflash_compression(
             http_detail::canonicalize_pflash_token_spans(std::move(pins));
         unified_request.query_suffix_candidates = query_suffix_candidates;
         unified_request.history_query_spans = history_query_spans;
+        unified_request.keep_query_spans = keep_query_spans;
         unified_request.turn_query_span = turn_query_span;
         unified_request.extra_query_spans = extra_query_spans;
         unified_request.request_messages = request_probe_messages;
@@ -6306,6 +6326,7 @@ std::string HttpServer::apply_pflash_compression(
         std::move(required_instruction_spans);
     compress_request.query_suffix_candidates = query_suffix_candidates;
     compress_request.history_query_spans = history_query_spans;
+    compress_request.keep_query_spans = keep_query_spans;
     compress_request.turn_query_span = turn_query_span;
     compress_request.extra_query_spans = extra_query_spans;
     compress_request.request_messages = request_probe_messages;

server/test/pflash/test_pflash_scoring_resume.cpp

@@ -123,6 +123,22 @@ TEST_CASE(PFlashScoringResumeFixture, query_sets_store_exact_rows_only) {
     CHECK(sets.size() == 0);
 }
 
+// Query windows are row sets keyed by their positions: a window stored before
+// a divergence point stays, one that reaches it is dropped.
+TEST_CASE(PFlashScoringResumeFixture, query_windows_are_position_sets) {
+    CHECK(pflash_span_tokens(5, 8) == std::vector<int>({5, 6, 7}));
+    CHECK(pflash_span_tokens(8, 8).empty());
+    PFlashExactRowSets windows;
+    const int hidden = 1;
+    const float row[] = {1};
+    std::vector<float> out;
+    CHECK(windows.put(pflash_span_tokens(447, 448), row, hidden, 65536));
+    CHECK(windows.put(pflash_span_tokens(12000, 12001), row, hidden, 65536));
+    CHECK(windows.invalidate_from(11000) == 1);
+    CHECK(windows.lookup(pflash_span_tokens(447, 448), hidden, out));
+    CHECK(!windows.lookup(pflash_span_tokens(12000, 12001), hidden, out));
+}
+
 // Session capacity: headroom S + S/2 + 4096, clamped to the context limit, never below S padded to 256.
 TEST_CASE(PFlashScoringResumeFixture, session_capacity_is_headroom_clamped_to_the_context_limit) {
     using luce::pflash::pflash_session_capacity;

server/test/unit/test_server_unit.cpp

@@ -11368,9 +11368,85 @@ TEST_CASE(ServerUnitFixture,
     TEST_ASSERT_MSG(tokenizer.decode({ids.begin() + history.end,
                                       ids.begin() + history.end + 2}) == "Sure.",
                     text_of(history));
+    // A chat whose prompt ends at the latest question's reply has nothing to
+    // keep: that question's window is the query itself.
+    TEST_ASSERT(request.keep_query_spans.empty());
     unlink(path.c_str());
 }
 
+// A tool session's prompt ends after tool output: the latest question's
+// history window (its reply's header) is kept for the next request, which
+// scores exactly that window once a new question arrives.
+TEST_CASE(ServerUnitFixture, test_pflash_tool_session_keeps_next_history_query) {
+    luce_test::ScopedEnvVar mode{"PFLASH_SELECT_MODE", "budget_only"};
+    luce_test::ScopedEnvVar skeleton{"PFLASH_CHAT_SKELETON_TOKENS", nullptr};
+
+    std::string material;
+    for (int i = 0; i < 80; ++i) material += "filler words here. ";
+    std::vector<ChatMessage> messages{
+        {"system", "You are helpful.", ""},
+        {"user", "What is the first answer?", ""},
+        {"assistant", "Sure.", ""},
+        {"tool", "Here is text: " + material, ""},
+    };
+    const auto run = [&] (const std::vector<ChatMessage> & chat,
+                          ModelBackend::CompressRequest & out,
+                          std::vector<int32_t> & ids_out) {
+        const std::string rendered = render_chat_template(
+            chat, ChatFormat::QWEN3, /*add_generation_prompt=*/true,
+            /*enable_thinking=*/true);
+        const std::string path = write_pflash_bpe_tokenizer_fixture(
+            {"What", " is", " the", " answer", " first", "?", "user", "assistant",
+             "system", "tool", "\n", "Sure", ".", " filler", " words", " here",
+             "Here", " text", ":", " Next", " question"},
+            rendered);
+        Tokenizer tokenizer;
+        TEST_ASSERT(tokenizer.load_from_gguf(path.c_str()));
+        auto backend_owner = std::make_unique<MockPflashCompressBackend>();
+        MockPflashCompressBackend & backend = *backend_owner;
+        LuceEngine engine(std::move(backend_owner));
+        ServerConfig config;
+        config.pflash_keep_ratio = 1.0f;
+        config.prefix_cache_cap = 0;
+        config.prefill_cache_cap = 0;
+        {
+            HttpServer server(engine, tokenizer, config);
+            server.set_drafter_tokenizer(&tokenizer);
+            ParsedRequest request;
+            request.format = ApiFormat::OPENAI_CHAT;
+            json wire = json::array();
+            for (const auto & message : chat) {
+                wire.push_back({{"role", message.role}, {"content", message.content}});
+            }
+            request.messages = wire;
+            request.prompt_tokens = tokenizer.encode(rendered);
+            const std::string error =
+                HttpServerTestAccess::apply_pflash_compression(server, request);
+            TEST_ASSERT_MSG(error.empty(), error);
+        }
+        TEST_ASSERT(backend.compress_calls == 1);
+        out = backend.last_request;
+        ids_out = out.input_ids;
+        unlink(path.c_str());
+    };
+    ModelBackend::CompressRequest first;
+    std::vector<int32_t> first_ids;
+    run(messages, first, first_ids);
+    TEST_ASSERT(first.history_query_spans.empty());
+    TEST_ASSERT(first.keep_query_spans.size() == 1);
+
+    messages.push_back({"assistant", "Sure.", ""});
+    messages.push_back({"user", "Next question?", ""});
+    ModelBackend::CompressRequest second;
+    std::vector<int32_t> second_ids;
+    run(messages, second, second_ids);
+    TEST_ASSERT(second.history_query_spans.size() == 1);
+    const auto kept = first.keep_query_spans[0];
+    const auto history = second.history_query_spans[0];
+    TEST_ASSERT(kept.begin == history.begin && kept.end == history.end);
+    TEST_ASSERT(std::equal(first_ids.begin(), first_ids.begin() + kept.end, second_ids.begin()));
+}
+
 // Two turns over one document through prepare_prompt; returns the served
 // prompts, the view modes and how often the compressor ran.
 struct PflashTwoTurnRun {