diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 0da29ea33..2697741ab 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -342,6 +342,7 @@ add_library(dflash_common STATIC src/server/tool_memory.cpp src/server/sse_emitter.cpp src/server/prefix_cache.cpp + src/server/pin_friendly_prompt.cpp src/server/disk_prefix_cache.cpp src/server/freeze_history.cpp # ── Jinja chat-template engine (vendored under deps/llama.cpp/common/) ── diff --git a/server/docs/PIN_FRIENDLY_PROMPT.md b/server/docs/PIN_FRIENDLY_PROMPT.md new file mode 100644 index 000000000..b9241672c --- /dev/null +++ b/server/docs/PIN_FRIENDLY_PROMPT.md @@ -0,0 +1,218 @@ +# DiffPin + +**Making agent prompts prefix-cacheable by floating volatility out of the stable head** + +--- + +## In one sentence + +**DiffPin** finds the small changing part of a long prompt (a session clock, a banner, a one-line header), moves it *after* everything that stayed the same, and tells the inference engine: *pin the KV cache here — this contiguous block will come back*. + +--- + +## Why this exists + +Modern agents send a huge head on every turn: identity, rules, dozens of tool schemas. That head is expensive to **prefill** — the model must write a **KV cache** (the memory of “what I already saw”) before it can answer. + +Most of that head does not change. A few tokens do — often a line like *Conversation started: Friday, July 31…*. + +Prefix caching only works when the **beginning** of the prompt matches exactly. One changed date in the middle of the head, and the engine treats the whole head as new. You pay the long prefill again. + +```mermaid +flowchart LR + subgraph cold ["Without DiffPin"] + A1["Same tools + identity"] + A2["Different clock"] + A3["Full head looks new"] + A4["Rebuild all KV 💸"] + A1 --> A2 --> A3 --> A4 + end +``` + +The waste is not “bad networking.” It is **throwing away valid KV** because a tiny volatile island sat where the pin key could not ignore it. + +--- + +## The invention: DiffPin + +**Name:** DiffPin +**Idea:** Treat volatility as a *diff hunk*, not as fate. + +Compare today’s tools/system head to recent traffic. The shared parts form a **prefix** and a **suffix**. The disagreement is a small **middle** (the clock). DiffPin rewrites the head so stable tokens become one contiguous block, with the volatile middle after them: + +```text +Before: [ stable … TIME … stable … ] +After: [ stable ……… stable ][ TIME ][ end ] + ↑ pin here +``` + +```mermaid +flowchart TB + subgraph before ["What the client sent"] + B1["████████"] + B2["🕐 clock"] + B3["████"] + B1 --- B2 --- B3 + end + + subgraph after ["What DiffPin serves"] + A1["████████████"] + A2["🕐 clock"] + A1 --- A2 + end + + before -->|"diff → float the clock"| after + A1 -.->|"protected pin
reuse this KV"| KV["KV cache"] +``` + +Same information. Better *shape* for caching. + +--- + +## How to picture it + +Think of a highlighter and a sticky note. + +1. **Diff** — lay two prompts side by side; highlight what matches from the left and from the right; the unhighlighted island is the sticky note (time, model line, …). +2. **Float** — peel the sticky note off the middle and place it after the highlighted paper. +3. **Pin** — stack the highlighted paper in the KV closet. Next time the sticky note says a different day, the paper still matches. + +```mermaid +sequenceDiagram + participant Turn1 as Turn with July 30 + participant DiffPin + participant Turn2 as Turn with July 31 + participant KV as KV pin + + Turn1->>DiffPin: Long head + clock A + DiffPin->>KV: Remember contiguous stable head + Turn2->>DiffPin: Same head + clock B + DiffPin->>DiffPin: Diff spots only the clock + DiffPin->>DiffPin: Float clock after stable block + DiffPin->>KV: Stable block matches — restore + Note over Turn2,KV: Prefill only clock + new user text +``` + +--- + +## What “pin-friendly” means + +A layout is **pin-friendly** when the tokens you want to reuse form one uninterrupted prefix: + +```mermaid +flowchart LR + subgraph friendly ["Pin-friendly"] + S["STABLE STABLE STABLE"] + V["volatile"] + T["transcript…"] + S --> V --> T + end +``` + +Not pin-friendly when volatility punches a hole in the middle — the shared suffix cannot join the pin, because prefix caches only care about *beginnings*: + +```mermaid +flowchart LR + subgraph unfriendly ["Not pin-friendly"] + S1["STABLE"] + V["volatile"] + S2["STABLE"] + T["transcript…"] + S1 --> V --> S2 --> T + end +``` + +DiffPin’s job is to turn the second shape into the first. + +--- + +## Where it sits in the story + +```mermaid +flowchart TB + Agent["Agent / harness
sends messages + tools"] + Render["Chat template
→ token stream"] + DiffPin["DiffPin
diff · float · pin"] + Engine["Inference engine
restore KV · prefill suffix · decode"] + + Agent --> Render --> DiffPin --> Engine +``` + +DiffPin does not invent a new chat API. It does not need a conversation-id header. It works from **tokens the engine already has**, plus a short memory of recent tool-bearing heads. + +--- + +## What it is careful about + +DiffPin only floats a middle hunk when that hunk looks like **ephemeral noise** (small). If the tools themselves changed, the “middle” is huge — DiffPin leaves the prompt alone. Wrong merges would be worse than a cold prefill. + +```mermaid +flowchart TD + D{Diff middle size?} + D -->|"Small — a clock"| Float["Float it · pin the rest"] + D -->|"Huge — tools changed"| Leave["Leave alone · new pin if needed"] +``` + +End-of-message markers (the “this turn is over” tokens) stay at the end of the head so the chat shape still reads as a coherent system turn. + +The rewrite head stops at the **first end-of-message marker**, not at the PrefixCache chat boundary. Those boundaries sit *after* the next role-start; cutting there would let a floated middle cross into the user turn. With no chat boundaries (custom templates), DiffPin does not rewrite. + +--- + +## What you should feel in production + +| Moment | Without DiffPin | With DiffPin | +|---|---|---| +| First tool-heavy turn | Pay for the long head | Same — fill the pin once | +| Next turn, same tools, new clock | Pay for the long head again | Restore pin; prefill clock + new text | +| Multi-chat noise | Long head may get evicted | Protected pin prefers to stay | + +The win is wall-clock and GPU time: tens of seconds of head prefill collapsing to a short suffix prefill when the stable block hits. + +--- + +## Relation to other ideas + +```mermaid +mindmap + root((Reuse agent KV)) + DiffPin + Diff finds volatility + Float to make a contiguous pin + No special client headers + Harness layout + Put clocks in a later message + Ideal when clients cooperate + Slot / protect policy + Keep the pin from being thrashing away + Complements DiffPin +``` + +DiffPin is the engine-side invention when clients bury clocks inside one big system blob. A disciplined harness still helps; DiffPin is the safety net that makes the common messy case pin-friendly anyway. + +--- + +## Naming + +| Term | Meaning | +|---|---| +| **DiffPin** | The invention: diff → float volatility → pin contiguous stable prefix | +| **Pin** | The KV snapshot of that stable prefix | +| **Float** | Moving the volatile hunk after the stable block | +| **Pin-friendly** | Stable tokens form one uninterrupted prefix | + +--- + +## Closing picture + +```mermaid +flowchart TB + Q["Question: why did warm cache miss
when only the date changed?"] + A["Answer: the date sat inside the pin key"] + I["Invention: DiffPin"] + M["Method: diff the head, float the date,
pin the paper that stayed the same"] + + Q --> A --> I --> M +``` + +Agents will keep sending large, almost-stable prompts. DiffPin makes “almost” good enough for KV reuse — by shaping the prompt so the cache can see the stability that was there all along. diff --git a/server/docs/TOOL_PREFIX_CACHE.md b/server/docs/TOOL_PREFIX_CACHE.md index dc1e31740..7a71604aa 100644 --- a/server/docs/TOOL_PREFIX_CACHE.md +++ b/server/docs/TOOL_PREFIX_CACHE.md @@ -10,9 +10,15 @@ For Qwen chat templates, the first request is effectively: ```text [system + tool schemas] [user request] [assistant start] - ^ native snapshot boundary + ^ protected tools-boundary snapshot ``` +Turn 1 pays the tool-schema prefill once and commits a **protected** native +inline snapshot at the system/tools head (first chat boundary). Later turns +restore that head and prefill only the new conversation suffix. Progressive +deeper conversation boundaries remain unprotected LRU leaves so multi-chat +traffic cannot thrash the shared tools pin away. + The snapshot is a normal, backend-owned prefix snapshot. This is important for hybrid architectures such as Qwen3.5/3.6: it contains attention KV, recurrent state, convolution state, and the last-token seed together. Composing a @@ -21,9 +27,12 @@ a valid restore. ## Runtime behavior -- Turn 1 pays the tool-schema prefill once and commits a native inline snapshot. -- Turn 2 restores the snapshot keyed by the system/tool boundary. +- Turn 1 (tools present, tools head miss): pay the tool-schema prefill once and + commit a protected inline snapshot at the first chat boundary. +- Turn 2 restores the tools-head snapshot and can deepen to a later turn + boundary after the head is already restored. - Later turns can restore progressively deeper conversation boundaries. +- Protected tools pins are skipped by eviction while unprotected leaves exist. - The cache key is the exact token prefix. Changing a tool name, description, parameter, system prompt, or template produces a miss; incompatible KV state is never reused. @@ -73,3 +82,13 @@ For an A/B against an unmodified `main` server, use the same command with `--timing-only`. That mode accepts the older three-field `usage.timings` shape and reports prefill speedup, but deliberately does not claim cache correctness because `main` cannot expose the restored-token counts. + +## See also + +When volatile text (for example a session clock) sits *inside* the first system +message, the first chat boundary cannot exclude it. **DiffPin** can diff the +head, float that volatility after the stable block, and pin the contiguous +prefix. See [PIN_FRIENDLY_PROMPT.md](./PIN_FRIENDLY_PROMPT.md). `DFLASH_PPP` +(on by default) enables LCP pin-end annotation + sticky protect; +`DFLASH_PPP_REARRANGE=1` opts into the token-level float rewrite (off by +default — unconstrained peels can scramble tool JSON). diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index 46081fa23..60d62cc66 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -21,6 +21,7 @@ #include "sse_emitter.h" #include "prompt_normalize.h" #include "tool_hint.h" +#include "pin_friendly_prompt.h" #include "common/sha1.h" #include "freeze_history.h" @@ -399,6 +400,10 @@ int resolve_max_output_tokens(const json & body, int default_max_tokens) { return default_max_tokens; } +bool ppp_prefers_tools_boundary(bool ppp_enabled, bool has_tools) { + return ppp_enabled && has_tools; +} + // Sampler parameters. When the request omits a value, fall back to the // model card's sampling defaults (spec §3.3); when the card doesn't // supply one either, use the hard-coded default. @@ -953,6 +958,36 @@ HttpServer::HttpServer(ModelBackend & backend, } disk_cache_.init(); status_html_path_ = resolve_status_html(); + + // PPP env overrides (operator-facing; no CLI flags required). + auto env_truthy = [](const char * v) -> bool { + if (!v || !*v) return false; + return !(v[0] == '0' && v[1] == '\0') && + !(v[0] == 'f' || v[0] == 'F' || v[0] == 'n' || v[0] == 'N'); + }; + if (const char * e = std::getenv("DFLASH_PPP")) { + config_.ppp_enabled = env_truthy(e); + } + if (const char * e = std::getenv("DFLASH_PPP_REARRANGE")) { + config_.ppp_rearrange = env_truthy(e); + } + if (const char * e = std::getenv("DFLASH_PPP_LCP_WINDOW")) { + const int n = std::atoi(e); + if (n > 0) config_.ppp_lcp_window = n; + } + if (const char * e = std::getenv("DFLASH_PPP_MIN_PIN_TOKENS")) { + const int n = std::atoi(e); + if (n > 0) config_.ppp_min_pin_tokens = n; + } + if (const char * e = std::getenv("DFLASH_PPP_MAX_EPHEMERAL")) { + const int n = std::atoi(e); + if (n > 0) config_.ppp_max_ephemeral_tokens = n; + } + std::fprintf(stderr, + "[ppp] enabled=%d rearrange=%d lcp_window=%d min_pin=%d max_ephemeral=%d\n", + (int)config_.ppp_enabled, (int)config_.ppp_rearrange, + config_.ppp_lcp_window, config_.ppp_min_pin_tokens, + config_.ppp_max_ephemeral_tokens); } // Resolve path to share/status.html at startup. @@ -1837,7 +1872,32 @@ bool HttpServer::route_request(SocketHandle fd, const HttpRequest & hr) { apply_request_reasoning(body, req); // Bandit: parse session_id from extra_body (opt-in adaptive keep_ratio). req.session_id = parse_session_id_from_body(body); - if (!render_and_tokenize_request(fd, chat_messages, req)) return true; + + // PPP rearrange (optional): peel ephemeral system banners into a + // following system message so the first chat boundary is stable. + std::vector render_messages = chat_messages; + if (config_.ppp_enabled && config_.ppp_rearrange && !req.tools.empty()) { + auto layout = PinFriendlyPrompt::rearrange(chat_messages, true); + if (layout.rearranged) { + render_messages = std::move(layout.messages); + // Keep FlowKV / response formatting aligned with the served + // layout (FlowKV re-renders from req.messages). + if (req.messages.is_array() && !req.messages.empty() && + req.messages[0].value("role", "") == "system" && + render_messages.size() >= 2) { + req.messages[0]["content"] = render_messages[0].content; + json meta = { + {"role", "system"}, + {"content", render_messages[1].content}, + }; + req.messages.insert(req.messages.begin() + 1, std::move(meta)); + } + std::fprintf(stderr, + "[ppp] rearranged: peeled ephemeral system tail\n"); + } + } + + if (!render_and_tokenize_request(fd, render_messages, req)) return true; // count_tokens: short-circuit after tokenization. Skip generation // entirely — Anthropic's contract is just {"input_tokens": N}. @@ -2659,16 +2719,80 @@ HttpServer::GenerationCacheState HttpServer::prepare_generation_cache( const ParsedRequest & req, PreparedPrompt & prepared, GenerateRequest & generate_request) { auto & effective_prompt = prepared.tokens; + // Tool-heavy requests prefer the reusable system/tool boundary under eviction. + const bool prefer_inline_snap = !req.tools.empty(); + const bool prefer_tools_boundary = + ppp_prefers_tools_boundary(config_.ppp_enabled, prefer_inline_snap); + int forced_cut = req.pin_end_token; + + // PPP runs *before* lookup. Default (rearrange=0): annotate a sticky + // pin_end only — never mutate tokens. Token-level DiffPin rewrite + // (prefix|suffix|middle float) is opt-in via DFLASH_PPP_REARRANGE=1; + // unconstrained middle peels can scramble tool-schema JSON and yield + // empty post-tool completions. + bool ppp_rewrote = false; + if (config_.ppp_enabled && prefer_tools_boundary) { + const auto boundaries = find_all_boundaries( + effective_prompt, prefix_cache_.chat_markers()); + if (config_.ppp_rearrange) { + auto rewrite = PinFriendlyPrompt::diff_make_pin_friendly( + effective_prompt, boundaries, recent_tool_prefixes_, + prefix_cache_.chat_markers(), + config_.ppp_lcp_window, config_.ppp_min_pin_tokens, + config_.ppp_max_ephemeral_tokens); + if (rewrite.rewritten) { + effective_prompt = std::move(rewrite.tokens); + generate_request.prompt = effective_prompt; + ppp_rewrote = true; + // Full-cache hits/keys from unrearranged tokens are stale. + prepared.full_cache_hit_slot = -1; + prepared.full_cache_hit_len = 0; + prepared.full_cache_served_tokens = -1; + std::fprintf(stderr, + "[ppp] diff-rewrite prefix=%d suffix=%d middle=%d " + "pin_end=%d prompt=%zu\n", + rewrite.prefix_len, rewrite.suffix_len, rewrite.middle_len, + rewrite.pin_end, effective_prompt.size()); + } + if (forced_cut <= 0) forced_cut = rewrite.pin_end; + if (forced_cut > 0 && !rewrite.rewritten) { + std::fprintf(stderr, + "[ppp] pin_end=%d (no rewrite; prompt=%zu)\n", + forced_cut, effective_prompt.size()); + } + } else if (forced_cut <= 0) { + forced_cut = PinFriendlyPrompt::annotate_pin_end( + effective_prompt, boundaries, recent_tool_prefixes_, + config_.ppp_lcp_window, config_.ppp_min_pin_tokens); + if (forced_cut > 0) { + std::fprintf(stderr, + "[ppp] pin_end=%d (pin-only; prompt=%zu)\n", + forced_cut, effective_prompt.size()); + } + } + const auto remember_bounds = find_all_boundaries( + effective_prompt, prefix_cache_.chat_markers()); + const int remember_n = !remember_bounds.empty() + ? remember_bounds.front() + : (int)effective_prompt.size(); + PinFriendlyPrompt::remember_tool_prefix( + recent_tool_prefixes_, effective_prompt, remember_n, + config_.ppp_lcp_window); + } + GenerationCacheState cache; cache.cache_slot = prepared.full_cache_hit_slot; cache.prefix_len = prepared.full_cache_hit_len; cache.using_restore = cache.cache_slot >= 0; cache.disk_policy = req.disk_cache_policy; + cache.full_snap_key_effective = ppp_rewrote; - // Exact raw-prompt snapshots take priority over inline turn boundaries. + // Exact full-prompt snapshots. After a DiffPin rewrite, key by the + // tokens we actually serve (effective_prompt), not the client wire form. if (!cache.using_restore) { - auto [full_slot, full_len] = - prefix_cache_.lookup_full(req.prompt_tokens); + const auto & full_key = + ppp_rewrote ? effective_prompt : req.prompt_tokens; + auto [full_slot, full_len] = prefix_cache_.lookup_full(full_key); if (full_slot >= 0) { cache.cache_slot = full_slot; cache.prefix_len = full_len; @@ -2907,17 +3031,19 @@ HttpServer::GenerationCacheState HttpServer::prepare_generation_cache( // A generation can save only one snapshot during prefill. Tool-heavy // requests prefer the reusable system/tool boundary; otherwise an // enabled exact full-prompt cache retains its existing priority. - const bool prefer_inline_snap = !req.tools.empty(); auto prepare_inline = [&]() { const auto prepared_snapshot = prefix_cache_.prepare_inline_snap( effective_prompt, - cache.using_restore ? cache.prefix_len : 0); + cache.using_restore ? cache.prefix_len : 0, + prefer_tools_boundary, + forced_cut); cache.snap_slot = prepared_snapshot.first; cache.snap_cut = prepared_snapshot.second; }; auto prepare_full = [&]() { - cache.full_snap_slot = - prefix_cache_.prepare_full_snap(req.prompt_tokens); + const auto & full_key = cache.full_snap_key_effective + ? effective_prompt : req.prompt_tokens; + cache.full_snap_slot = prefix_cache_.prepare_full_snap(full_key); if (cache.full_snap_slot >= 0) { cache.full_snap_pos = (int) effective_prompt.size(); generate_request.snap_slot = cache.full_snap_slot; @@ -2994,8 +3120,10 @@ void HttpServer::finalize_generation_cache( backend_.snapshot_cur_pos(cache.full_snap_slot); if (saved_position > 0 && saved_position <= cache.full_snap_pos) { + const auto & full_key = cache.full_snap_key_effective + ? effective_prompt : req.prompt_tokens; prefix_cache_.confirm_full_snap( - cache.full_snap_slot, req.prompt_tokens, saved_position); + cache.full_snap_slot, full_key, saved_position); } else { backend_.snapshot_free(cache.full_snap_slot); prefix_cache_.abort_full_snap(cache.full_snap_slot); diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index 56f555369..dc6654ff5 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -61,6 +61,16 @@ struct ServerConfig { int prefix_cache_cap = 32; // prefix cache slots (0 disables) int prefill_cache_cap = 0; // full-prompt/prefill cache slots (0 disables) + // Pin-Friendly Prompt Processor (PPP): LCP pin_end + optional rearrange. + // See docs/PIN_FRIENDLY_PROMPT.md. Env: DFLASH_PPP=0|1, + // DFLASH_PPP_REARRANGE=0|1, DFLASH_PPP_LCP_WINDOW=N, + // DFLASH_PPP_MIN_PIN_TOKENS=N, DFLASH_PPP_MAX_EPHEMERAL=N. + bool ppp_enabled = true; + bool ppp_rearrange = false; + int ppp_lcp_window = 8; + int ppp_min_pin_tokens = 512; + int ppp_max_ephemeral_tokens = 256; // diff hunk relocate cap + // Thinking-budget v2. Applied when a request opts in via // `thinking: {type: "enabled"}` or `reasoning: {effort: ...}`. // think_max_tokens caps phase-1 reasoning generation; the combined @@ -228,6 +238,8 @@ struct ParsedRequest { // Bandit: per-session adaptive keep_ratio opt-in std::string session_id; DiskPrefixCachePolicy disk_cache_policy; + // PPP: stable pin cut for tool-heavy requests (0 = use default boundary). + int pin_end_token = 0; }; // Parse request sampler fields, applying model-card defaults where present. @@ -243,6 +255,10 @@ json require_messages_array(const json & body); // selected field is parsed, so malformed lower-priority aliases are ignored. int resolve_max_output_tokens(const json & body, int default_max_tokens); +// Sticky tools-boundary pinning is part of PPP and must follow its master +// toggle. Kept as a small policy helper so the disabled path is testable. +bool ppp_prefers_tools_boundary(bool ppp_enabled, bool has_tools); + // Build the /props response body. Exposed (non-static) so unit tests // can assert on its shape without spinning up a real socket. See // docs/specs/props-endpoint.md for the wire contract. @@ -320,6 +336,9 @@ class HttpServer { int full_snap_slot = -1; int full_snap_pos = 0; bool full_snap_prepared = false; + // When DiffPin rewrote tokens, full-cache keys must use + // prepared.tokens (effective), not req.prompt_tokens. + bool full_snap_key_effective = false; int snap_slot = -1; int snap_cut = 0; bool snap_prepared = false; @@ -433,6 +452,8 @@ class HttpServer { // Track prompt tokens for each snapshot slot (for shutdown save). std::unordered_map> slot_tokens_; std::vector> recent_disk_prompts_; + // Recent tool-bearing prompt prefixes for PPP LCP annotate. + std::vector> recent_tool_prefixes_; // FlowKV freeze-history: per-message compression cache. // Key: SHA-1 hash of the drafter-token slice for an aged message. diff --git a/server/src/server/pin_friendly_prompt.cpp b/server/src/server/pin_friendly_prompt.cpp new file mode 100644 index 000000000..82b6fc2c0 --- /dev/null +++ b/server/src/server/pin_friendly_prompt.cpp @@ -0,0 +1,323 @@ +#include "pin_friendly_prompt.h" + +#include +#include + +namespace dflash::common { + +int PinFriendlyPrompt::longest_common_prefix_len( + const std::vector & a, + const std::vector & b) { + const int n = (int)std::min(a.size(), b.size()); + int i = 0; + while (i < n && a[(size_t)i] == b[(size_t)i]) ++i; + return i; +} + +int PinFriendlyPrompt::longest_common_suffix_len( + const std::vector & a, + const std::vector & b, + int prefix_len) { + int i = (int)a.size() - 1; + int j = (int)b.size() - 1; + int s = 0; + while (i >= prefix_len && j >= prefix_len && + a[(size_t)i] == b[(size_t)j]) { + --i; + --j; + ++s; + } + return s; +} + +TokenDiffSplit PinFriendlyPrompt::diff_split( + const std::vector & reference, + const std::vector & current) { + TokenDiffSplit out; + out.prefix_len = longest_common_prefix_len(reference, current); + out.suffix_len = longest_common_suffix_len( + reference, current, out.prefix_len); + out.middle_begin = out.prefix_len; + out.middle_end = (int)current.size() - out.suffix_len; + if (out.middle_end < out.middle_begin) { + out.middle_end = out.middle_begin; + out.suffix_len = (int)current.size() - out.prefix_len; + } + return out; +} + +int PinFriendlyPrompt::safe_boundary_cut( + int n, const std::vector & boundaries) { + if (n <= 0) return 0; + int best = 0; + for (int b : boundaries) { + if (b > 0 && b <= n) best = std::max(best, b); + } + return best; +} + +int PinFriendlyPrompt::choose_pin_end( + int lcp, + const std::vector & boundaries, + int min_pin_tokens) { + if (lcp < min_pin_tokens) return 0; + const int boundary = safe_boundary_cut(lcp, boundaries); + if (boundary >= min_pin_tokens) return boundary; + return lcp; +} + +int PinFriendlyPrompt::annotate_pin_end( + const std::vector & tokens, + const std::vector & boundaries, + const std::vector> & recent_tool_prefixes, + int window, + int min_pin_tokens) { + if (tokens.empty() || recent_tool_prefixes.empty() || window <= 0) { + return 0; + } + const int n = std::min(window, (int)recent_tool_prefixes.size()); + int common = 0; + for (int i = 0; i < n; ++i) { + const int idx = (int)recent_tool_prefixes.size() - n + i; + common = std::max( + common, + longest_common_prefix_len(tokens, recent_tool_prefixes[(size_t)idx])); + } + return choose_pin_end(common, boundaries, min_pin_tokens); +} + +int PinFriendlyPrompt::trailing_end_marker_len( + const std::vector & ids, + const ChatMarkers & markers) { + if (ids.empty()) return 0; + int best = 0; + for (const auto & seq : markers.end_msg_seqs) { + if (seq.empty() || (int)seq.size() > (int)ids.size()) continue; + const int start = (int)ids.size() - (int)seq.size(); + bool match = true; + for (int i = 0; i < (int)seq.size(); ++i) { + if (ids[(size_t)(start + i)] != seq[(size_t)i]) { + match = false; + break; + } + } + if (match) best = std::max(best, (int)seq.size()); + } + // Qwen boundaries often include a trailing '\n' token after im_end. + // Peek one extra equal-length attempt is unnecessary; keep marker only. + return best; +} + +int PinFriendlyPrompt::tools_system_head_end( + const std::vector & tokens, + const ChatMarkers & markers) { + if (tokens.empty() || markers.end_msg_seqs.empty()) return 0; + int best = -1; + int best_len = 0; + for (const auto & seq : markers.end_msg_seqs) { + if (seq.empty() || (int)seq.size() > (int)tokens.size()) continue; + for (int i = 0; i + (int)seq.size() <= (int)tokens.size(); ++i) { + bool match = true; + for (int k = 0; k < (int)seq.size(); ++k) { + if (tokens[(size_t)(i + k)] != seq[(size_t)k]) { + match = false; + break; + } + } + if (!match) continue; + if (best < 0 || i < best) { + best = i; + best_len = (int)seq.size(); + } + break; // first occurrence of this seq; keep scanning other seqs + } + } + if (best < 0) return 0; + return best + best_len; +} + +PinFriendlyRewrite PinFriendlyPrompt::diff_make_pin_friendly( + const std::vector & tokens, + const std::vector & boundaries, + const std::vector> & recent_tool_prefixes, + const ChatMarkers & markers, + int window, + int min_pin_tokens, + int max_ephemeral_tokens) { + PinFriendlyRewrite out; + out.tokens = tokens; + if (tokens.empty() || recent_tool_prefixes.empty() || window <= 0) { + return out; + } + // Custom / unstructured templates: do not rewrite the whole prompt. + if (boundaries.empty()) return out; + + // Only rewrite through the first end-of-message marker. find_all_boundaries + // returns cuts *after* the next role-start; using those would let the + // volatile middle float past user/assistant markers. + const int head_end = tools_system_head_end(tokens, markers); + if (head_end <= 0 || head_end < min_pin_tokens) return out; + + std::vector head(tokens.begin(), tokens.begin() + head_end); + const std::vector rest(tokens.begin() + head_end, tokens.end()); + + const int trailer_len = trailing_end_marker_len(head, markers); + std::vector trailer; + if (trailer_len > 0) { + trailer.assign(head.end() - trailer_len, head.end()); + head.resize((size_t)((int)head.size() - trailer_len)); + } + if ((int)head.size() < min_pin_tokens) return out; + + // Pick the reference that yields the largest stable span + // (prefix + suffix) with a small ephemeral middle. + TokenDiffSplit best{}; + int best_stable = -1; + bool found = false; + const int n = std::min(window, (int)recent_tool_prefixes.size()); + for (int i = 0; i < n; ++i) { + const int idx = (int)recent_tool_prefixes.size() - n + i; + const auto & ref_full = recent_tool_prefixes[(size_t)idx]; + // Align reference to body length (strip its own trailer if present). + std::vector ref = ref_full; + if ((int)ref.size() > head_end) { + ref.resize((size_t)head_end); + } + const int ref_trailer = trailing_end_marker_len(ref, markers); + if (ref_trailer > 0 && ref_trailer < (int)ref.size()) { + ref.resize((size_t)((int)ref.size() - ref_trailer)); + } + if (ref.empty()) continue; + + const TokenDiffSplit split = diff_split(ref, head); + const int middle = split.middle_end - split.middle_begin; + const int stable = split.prefix_len + split.suffix_len; + if (middle < 0 || middle > max_ephemeral_tokens) continue; + if (stable < min_pin_tokens) continue; + // Require a real middle relocation opportunity (suffix beyond prefix). + if (middle == 0) continue; + if (stable > best_stable) { + best_stable = stable; + best = split; + found = true; + } + } + + if (!found) { + // Fall back: LCP pin without rewrite. + out.pin_end = annotate_pin_end( + tokens, boundaries, recent_tool_prefixes, window, min_pin_tokens); + return out; + } + + std::vector new_head; + new_head.reserve(head.size() + trailer.size()); + new_head.insert(new_head.end(), + head.begin(), head.begin() + best.prefix_len); + if (best.suffix_len > 0) { + new_head.insert(new_head.end(), + head.end() - best.suffix_len, head.end()); + } + const int pin_body = (int)new_head.size(); + new_head.insert(new_head.end(), + head.begin() + best.middle_begin, + head.begin() + best.middle_end); + new_head.insert(new_head.end(), trailer.begin(), trailer.end()); + + out.tokens.clear(); + out.tokens.insert(out.tokens.end(), new_head.begin(), new_head.end()); + out.tokens.insert(out.tokens.end(), rest.begin(), rest.end()); + out.pin_end = pin_body; // contiguous stable blob; volatile + trailer after + // Prefer including trailer in the pin when it yields a chat boundary cut. + if (!trailer.empty()) { + const int with_trailer = pin_body + (int)trailer.size(); + // Only if volatile sits after trailer we would have moved wrong; + // here volatile is before trailer, so pin stays at pin_body. + (void)with_trailer; + } + out.rewritten = (out.tokens != tokens); + out.prefix_len = best.prefix_len; + out.suffix_len = best.suffix_len; + out.middle_len = best.middle_end - best.middle_begin; + + if (out.pin_end < min_pin_tokens) { + out.tokens = tokens; + out.rewritten = false; + out.pin_end = annotate_pin_end( + tokens, boundaries, recent_tool_prefixes, window, min_pin_tokens); + } + return out; +} + +static std::string trim_leading_newlines(std::string s) { + while (!s.empty() && (s[0] == '\n' || s[0] == '\r')) { + s.erase(s.begin()); + } + return s; +} + +std::pair +PinFriendlyPrompt::split_ephemeral_system_tail(const std::string & system) { + static const char * kMarkers[] = { + "\nConversation started:", + "\n\nConversation started:", + "\nSession started:", + "\n\nSession started:", + }; + size_t cut = std::string::npos; + for (const char * marker : kMarkers) { + const size_t p = system.rfind(marker); + if (p == std::string::npos) continue; + if (cut == std::string::npos || p < cut) cut = p; + } + if (cut == std::string::npos || cut == 0) { + return {system, {}}; + } + std::string stable = system.substr(0, cut); + while (!stable.empty() && + (stable.back() == ' ' || stable.back() == '\t' || + stable.back() == '\n' || stable.back() == '\r')) { + stable.pop_back(); + } + std::string ephemeral = trim_leading_newlines(system.substr(cut)); + if (ephemeral.empty()) return {system, {}}; + return {std::move(stable), std::move(ephemeral)}; +} + +PinFriendlyLayout PinFriendlyPrompt::rearrange( + const std::vector & messages, bool enable) { + PinFriendlyLayout layout; + layout.messages = messages; + if (!enable || messages.empty() || messages[0].role != "system") { + return layout; + } + auto [stable, ephemeral] = split_ephemeral_system_tail(messages[0].content); + if (ephemeral.empty() || stable.empty()) { + return layout; + } + layout.messages[0].content = std::move(stable); + ChatMessage meta; + meta.role = "system"; + meta.content = std::move(ephemeral); + layout.messages.insert(layout.messages.begin() + 1, std::move(meta)); + layout.rearranged = true; + return layout; +} + +void PinFriendlyPrompt::remember_tool_prefix( + std::vector> & ring, + const std::vector & prompt_ids, + int max_prefix_tokens, + int window) { + if (prompt_ids.empty() || window <= 0) return; + const int n = std::min((int)prompt_ids.size(), + max_prefix_tokens > 0 ? max_prefix_tokens + : (int)prompt_ids.size()); + std::vector prefix(prompt_ids.begin(), prompt_ids.begin() + n); + ring.push_back(std::move(prefix)); + while ((int)ring.size() > window) { + ring.erase(ring.begin()); + } +} + +} // namespace dflash::common diff --git a/server/src/server/pin_friendly_prompt.h b/server/src/server/pin_friendly_prompt.h new file mode 100644 index 000000000..37db62e4b --- /dev/null +++ b/server/src/server/pin_friendly_prompt.h @@ -0,0 +1,111 @@ +// Pin-Friendly Prompt Processor (PPP) +// +// Diffs the tools/system head against recent traffic, isolates the volatile +// span, and rewrites tokens to: +// [shared prefix][shared suffix][volatile middle][end markers] +// so PrefixCache can pin the contiguous stable blob. See docs/PIN_FRIENDLY_PROMPT.md. + +#pragma once + +#include "chat_template.h" +#include "prefix_cache.h" // ChatMarkers + +#include +#include +#include +#include + +namespace dflash::common { + +struct PinFriendlyLayout { + std::vector messages; + int pin_end_token = 0; + bool rearranged = false; + int lcp_len = 0; +}; + +// Result of a prefix/suffix diff on one sequence against a reference. +struct TokenDiffSplit { + int prefix_len = 0; // shared head + int suffix_len = 0; // shared tail (non-overlapping with prefix) + int middle_begin = 0; // in `current` + int middle_end = 0; // exclusive in `current` +}; + +// Full-prompt rewrite for pin-friendly layout. +struct PinFriendlyRewrite { + std::vector tokens; + int pin_end = 0; + bool rewritten = false; + int prefix_len = 0; + int suffix_len = 0; + int middle_len = 0; +}; + +class PinFriendlyPrompt { +public: + static int longest_common_prefix_len(const std::vector & a, + const std::vector & b); + + static int longest_common_suffix_len(const std::vector & a, + const std::vector & b, + int prefix_len); + + // Split `current` vs `reference` into shared prefix / volatile middle / + // shared suffix (classic single-hunk diff). + static TokenDiffSplit diff_split(const std::vector & reference, + const std::vector & current); + + // Max chat boundary ≤ n (0 if none). + static int safe_boundary_cut(int n, const std::vector & boundaries); + + static int choose_pin_end(int lcp, + const std::vector & boundaries, + int min_pin_tokens); + + // Legacy LCP-only annotate (no rewrite). + static int annotate_pin_end( + const std::vector & tokens, + const std::vector & boundaries, + const std::vector> & recent_tool_prefixes, + int window, + int min_pin_tokens); + + // Exclusive end of the first tools/system message (through end-of-message + // marker). Chat boundaries from find_all_boundaries() sit *after* the next + // role-start and must not be used as the DiffPin rewrite head. + static int tools_system_head_end(const std::vector & tokens, + const ChatMarkers & markers); + + // Diff the tools/system head against recent prefixes and rewrite: + // [tokentokentoken] with a mid clock → [tokentoken][token][time][im_end] + // pin_end covers the contiguous stable blob. No-op when boundaries are + // empty (custom templates), history is empty, the volatile hunk is too + // large, or nothing moves. + static PinFriendlyRewrite diff_make_pin_friendly( + const std::vector & tokens, + const std::vector & boundaries, + const std::vector> & recent_tool_prefixes, + const ChatMarkers & markers, + int window, + int min_pin_tokens, + int max_ephemeral_tokens); + + static std::pair + split_ephemeral_system_tail(const std::string & system); + + static PinFriendlyLayout rearrange(const std::vector & messages, + bool enable); + + static void remember_tool_prefix( + std::vector> & ring, + const std::vector & prompt_ids, + int max_prefix_tokens, + int window); + + // Peel trailing chat end-message marker tokens from `ids`. + static int trailing_end_marker_len(const std::vector & ids, + const ChatMarkers & markers); +}; + +} // namespace dflash::common diff --git a/server/src/server/prefix_cache.cpp b/server/src/server/prefix_cache.cpp index 4ab20c590..aa9a33136 100644 --- a/server/src/server/prefix_cache.cpp +++ b/server/src/server/prefix_cache.cpp @@ -149,32 +149,48 @@ static bool is_strict_prefix(const std::vector & a, return std::equal(a.begin(), a.end(), b.begin()); } -int select_inline_evict_victim(const std::vector *> & ids_lru) { +int select_inline_evict_victim(const std::vector *> & ids_lru, + const std::vector * protected_lru) { const int n = (int)ids_lru.size(); if (n <= 0) return 0; - // Oldest-first scan: evict the first entry that is not a strict prefix of any - // other entry (a leaf). Shared ancestor prefixes are thereby kept resident. + auto is_protected = [&](int i) { + return protected_lru && i >= 0 && i < (int)protected_lru->size() && + (*protected_lru)[(size_t)i]; + }; + // Oldest-first scan: prefer an unprotected leaf so sticky tools pins survive. + int oldest_protected_leaf = -1; for (int i = 0; i < n; i++) { bool is_ancestor = false; for (int j = 0; j < n; j++) { if (j == i) continue; if (is_strict_prefix(*ids_lru[i], *ids_lru[j])) { is_ancestor = true; break; } } - if (!is_ancestor) return i; // oldest leaf + if (is_ancestor) continue; + if (!is_protected(i)) return i; // oldest unprotected leaf + if (oldest_protected_leaf < 0) oldest_protected_leaf = i; } + if (oldest_protected_leaf >= 0) return oldest_protected_leaf; return 0; // unreachable (the longest entry is always a leaf); pure-LRU fallback } -int select_inline_evict_victim(const std::vector> & ids_lru) { +int select_inline_evict_victim(const std::vector> & ids_lru, + const std::vector * protected_lru) { std::vector *> ptrs; ptrs.reserve(ids_lru.size()); for (const auto & v : ids_lru) ptrs.push_back(&v); - return select_inline_evict_victim(ptrs); + return select_inline_evict_victim(ptrs, protected_lru); } int select_inline_snapshot_boundary(const std::vector & boundaries, - int restored_prefix_len) { + int restored_prefix_len, + bool prefer_tools_boundary) { if (boundaries.empty()) return 0; + // Tool-heavy cold path: pin the system+tools head (first marker) before + // deepening into conversation turns. Matches Python thin-pin semantics. + if (prefer_tools_boundary) { + const int tools_cut = boundaries.front(); + if (tools_cut > restored_prefix_len) return tools_cut; + } const int target = boundaries.size() >= 2 ? boundaries[boundaries.size() - 2] : boundaries.back(); @@ -238,6 +254,7 @@ std::pair PrefixCache::lookup(const std::vector & prompt_ids) auto boundaries = find_all_boundaries(prompt_ids, markers_); int best_slot = -1, best_len = 0; + int best_idx = -1; for (int cut : boundaries) { auto key = hash_prefix(prompt_ids.data(), cut); @@ -257,12 +274,27 @@ std::pair PrefixCache::lookup(const std::vector & prompt_ids) if (cut > best_len) { best_slot = entries_[idx].slot; best_len = cut; + best_idx = idx; } - move_to_end(idx); } } - if (best_slot >= 0) { + // Match committed entry prefixes directly. Required for PPP mid-message + // pin_end cuts that are not chat-template boundaries. + for (int i = 0; i < (int)entries_.size(); ++i) { + const auto & e = entries_[(size_t)i]; + const int len = (int)e.ids.size(); + if (len <= best_len || len > (int)prompt_ids.size()) continue; + if (!std::equal(e.ids.begin(), e.ids.end(), prompt_ids.begin())) { + continue; + } + best_slot = e.slot; + best_len = len; + best_idx = i; + } + + if (best_idx >= 0) { + move_to_end(best_idx); lifetime_hits_.fetch_add(1, std::memory_order_relaxed); std::fprintf(stderr, "[pc] lookup hit slot=%d prefix_len=%d (of %zu total)\n", best_slot, best_len, prompt_ids.size()); @@ -272,34 +304,57 @@ std::pair PrefixCache::lookup(const std::vector & prompt_ids) std::pair PrefixCache::prepare_inline_snap( const std::vector & prompt_ids, - int restored_prefix_len) { + int restored_prefix_len, + bool prefer_tools_boundary, + int forced_cut) { if (disabled_) return {-1, 0}; auto candidates = find_all_boundaries(prompt_ids, markers_); - const int target_cut = - select_inline_snapshot_boundary(candidates, restored_prefix_len); + int target_cut = 0; + bool forced = false; + if (forced_cut > restored_prefix_len && + forced_cut <= (int)prompt_ids.size()) { + target_cut = forced_cut; + forced = true; + } else { + target_cut = select_inline_snapshot_boundary( + candidates, restored_prefix_len, prefer_tools_boundary); + } if (target_cut <= 0) return {-1, 0}; auto key = hash_prefix(prompt_ids.data(), target_cut); if (find_entry(key) >= 0) return {-1, 0}; // already cached + // Protect the tools head pin for tool-heavy requests so multi-chat deepen + // snaps cannot thrash the ~18k system+tools KV away. PPP forced cuts are + // the stable tools/identity span and stay protected as well. + pending_protect_ = prefer_tools_boundary && + (forced || + (!candidates.empty() && target_cut == candidates.front())); + int slot; if ((int)entries_.size() >= cap_) { // At capacity — reserve a slot without evicting yet. Prefix-aware: prefer // the oldest leaf so shared ancestor prefixes (reused by later branches) - // stay resident. entries_ is already in LRU order (front = oldest). + // stay resident. Skip protected tools pins when an unprotected leaf exists. std::vector *> ids_lru; + std::vector protected_lru; ids_lru.reserve(entries_.size()); - for (const auto & e : entries_) ids_lru.push_back(&e.ids); - int victim = select_inline_evict_victim(ids_lru); + protected_lru.reserve(entries_.size()); + for (const auto & e : entries_) { + ids_lru.push_back(&e.ids); + protected_lru.push_back(e.protect); + } + int victim = select_inline_evict_victim(ids_lru, &protected_lru); pending_evict_key_ = entries_[victim].hash; has_pending_evict_ = true; slot = entries_[victim].slot; - if (victim != 0) { + if (victim != 0 || entries_[victim].protect) { std::fprintf(stderr, - "[pc] prefix-aware evict: victim idx=%d (len=%zu) kept oldest " - "ancestor (len=%zu)\n", - victim, entries_[victim].ids.size(), entries_.front().ids.size()); + "[pc] prefix-aware evict: victim idx=%d protect=%d (len=%zu) " + "kept oldest ancestor (len=%zu)\n", + victim, (int)entries_[victim].protect, + entries_[victim].ids.size(), entries_.front().ids.size()); } } else { slot = next_slot_; @@ -311,7 +366,8 @@ std::pair PrefixCache::prepare_inline_snap( } void PrefixCache::confirm_inline_snap(int slot, int target_cut, - const std::vector & prompt_ids) { + const std::vector & prompt_ids, + bool protect) { if (disabled_) return; // Evict the reserved entry (if any). @@ -339,12 +395,16 @@ void PrefixCache::confirm_inline_snap(int slot, int target_cut, } } + const bool protect_entry = protect || pending_protect_; + pending_protect_ = false; + auto key = hash_prefix(prompt_ids.data(), target_cut); std::vector ids(prompt_ids.begin(), prompt_ids.begin() + target_cut); - entries_.push_back({key, slot, std::move(ids)}); + entries_.push_back({key, slot, std::move(ids), protect_entry}); entries_size_count_.fetch_add(1, std::memory_order_relaxed); - std::fprintf(stderr, "[pc] inline-snap committed slot=%d prefix_len=%d\n", - slot, target_cut); + std::fprintf(stderr, + "[pc] inline-snap committed slot=%d prefix_len=%d protect=%d\n", + slot, target_cut, (int)protect_entry); } void PrefixCache::abort_inline_snap(int slot) { @@ -360,6 +420,7 @@ void PrefixCache::abort_inline_snap(int slot) { } } has_pending_evict_ = false; + pending_protect_ = false; } void PrefixCache::cancel_inline_snap(int slot) { @@ -369,6 +430,7 @@ void PrefixCache::cancel_inline_snap(int slot) { if (idx >= 0 && entries_[idx].slot != slot) return; } has_pending_evict_ = false; + pending_protect_ = false; } void PrefixCache::mark_all_cleared() { @@ -378,6 +440,7 @@ void PrefixCache::mark_all_cleared() { entries_size_count_.store(0, std::memory_order_relaxed); next_slot_ = 0; has_pending_evict_ = false; + pending_protect_ = false; std::fprintf(stderr, "[pc] all-cleared — dropped %d LRU entries\n", n); } diff --git a/server/src/server/prefix_cache.h b/server/src/server/prefix_cache.h index 2a0515749..a4683c69d 100644 --- a/server/src/server/prefix_cache.h +++ b/server/src/server/prefix_cache.h @@ -54,14 +54,25 @@ PrefixHash hash_prefix(const int32_t * ids, int count); // The pointer overload is the core (the caller passes pointers into its own // entries so no token vectors are copied); the value overload is a convenience // wrapper for tests. -int select_inline_evict_victim(const std::vector *> & ids_lru); -int select_inline_evict_victim(const std::vector> & ids_lru); - -// Pick the inline snapshot boundary for a request. We cache the boundary before -// the current user turn (second-to-last marker) and only when it advances past -// an already-restored prefix. Returns 0 when there is no useful new boundary. +// +// When `protected_lru` is non-null and same-sized, entries with +// `(*protected_lru)[i] == true` are skipped unless every leaf is protected +// (then the oldest protected leaf is chosen as a last resort). +int select_inline_evict_victim(const std::vector *> & ids_lru, + const std::vector * protected_lru = nullptr); +int select_inline_evict_victim(const std::vector> & ids_lru, + const std::vector * protected_lru = nullptr); + +// Pick the inline snapshot boundary for a request. +// Default: boundary before the current user turn (second-to-last marker), +// only when it advances past an already-restored prefix. +// When prefer_tools_boundary is set (tool-heavy agent requests), prefer the +// first marker (system+tools head) until that cut is already restored — this +// is the sticky "thin pin" Python tool-split used to keep under multi-chat +// eviction. Returns 0 when there is no useful new boundary. int select_inline_snapshot_boundary(const std::vector & boundaries, - int restored_prefix_len = 0); + int restored_prefix_len = 0, + bool prefer_tools_boundary = false); // ─── Prefix cache entry ───────────────────────────────────────────────── @@ -94,15 +105,22 @@ class PrefixCache { std::pair lookup(const std::vector & prompt_ids); // Prepare an inline snapshot. `restored_prefix_len` prevents reserving a - // slot for a boundary already covered by the restored snapshot. Returns - // (slot, target_cut) or (-1, 0). + // slot for a boundary already covered by the restored snapshot. + // `prefer_tools_boundary` selects the system/tools head first (see + // select_inline_snapshot_boundary). When `forced_cut` > restored, that + // cut is used instead (PPP pin_end, including mid-message LCP cuts). + // Returns (slot, target_cut) or (-1, 0). std::pair prepare_inline_snap( const std::vector & prompt_ids, - int restored_prefix_len = 0); + int restored_prefix_len = 0, + bool prefer_tools_boundary = false, + int forced_cut = 0); // Confirm after daemon successfully saved the snapshot. + // `protect` marks the entry non-evictable by unprotected traffic (tool pin). void confirm_inline_snap(int slot, int target_cut, - const std::vector & prompt_ids); + const std::vector & prompt_ids, + bool protect = false); // Abort if the snapshot failed. void abort_inline_snap(int slot); @@ -169,7 +187,10 @@ class PrefixCache { PrefixHash hash; int slot; std::vector ids; // prefix tokens [0, target_cut) for prefix-aware eviction + bool protect = false; // sticky tools-boundary pin }; + // Pending protect flag for the in-flight reservation (applied on confirm). + bool pending_protect_ = false; std::vector entries_; int next_slot_ = 0; PrefixHash pending_evict_key_{}; diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index 06e3521d4..68c07b703 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -12,6 +12,7 @@ #include "server/tool_parser.h" #include "server/reasoning.h" #include "server/prefix_cache.h" +#include "server/pin_friendly_prompt.h" #include "server/disk_prefix_cache.h" #include "server/freeze_history.h" #include "server/utf8_utils.h" @@ -1799,6 +1800,183 @@ TEST_CASE(ServerUnitFixture, test_inline_snapshot_boundary_advances_past_restore TEST_ASSERT(select_inline_snapshot_boundary({100}, 0) == 100); } +TEST_CASE(ServerUnitFixture, test_inline_snapshot_prefers_tools_boundary_until_restored) { + const std::vector boundaries = {100, 240, 380, 520}; + // Cold tool-heavy: pin system+tools head (first marker), not deepen cut. + TEST_ASSERT(select_inline_snapshot_boundary(boundaries, 0, true) == 100); + // After tools head is restored, deepen to second-to-last. + TEST_ASSERT(select_inline_snapshot_boundary(boundaries, 100, true) == 380); + TEST_ASSERT(select_inline_snapshot_boundary(boundaries, 380, true) == 0); + TEST_ASSERT(select_inline_snapshot_boundary({100}, 0, true) == 100); + TEST_ASSERT(select_inline_snapshot_boundary({100}, 100, true) == 0); +} + +TEST_CASE(ServerUnitFixture, test_ppp_master_toggle_gates_tools_boundary_pinning) { + TEST_ASSERT(ppp_prefers_tools_boundary(true, true)); + TEST_ASSERT(!ppp_prefers_tools_boundary(false, true)); + TEST_ASSERT(!ppp_prefers_tools_boundary(true, false)); + TEST_ASSERT(!ppp_prefers_tools_boundary(false, false)); +} + +// ── Pin-Friendly Prompt Processor (PPP) ───────────────────────────────── + +TEST_CASE(ServerUnitFixture, test_ppp_lcp_and_safe_boundary) { + const std::vector a = {1, 2, 3, 4, 5, 6}; + const std::vector b = {1, 2, 3, 9, 9}; + TEST_ASSERT(PinFriendlyPrompt::longest_common_prefix_len(a, b) == 3); + TEST_ASSERT(PinFriendlyPrompt::longest_common_prefix_len(a, a) == 6); + TEST_ASSERT(PinFriendlyPrompt::longest_common_prefix_len(a, {}) == 0); + + const std::vector boundaries = {100, 240, 380}; + TEST_ASSERT(PinFriendlyPrompt::safe_boundary_cut(250, boundaries) == 240); + TEST_ASSERT(PinFriendlyPrompt::safe_boundary_cut(50, boundaries) == 0); + TEST_ASSERT(PinFriendlyPrompt::safe_boundary_cut(380, boundaries) == 380); +} + +TEST_CASE(ServerUnitFixture, test_ppp_choose_pin_end_prefers_boundary_then_mid) { + const std::vector boundaries = {100, 200}; + // LCP past a boundary → pin at that boundary. + TEST_ASSERT(PinFriendlyPrompt::choose_pin_end(150, boundaries, 50) == 100); + // LCP past first boundary but short of second → still prefer boundary. + TEST_ASSERT(PinFriendlyPrompt::choose_pin_end(175, boundaries, 50) == 100); + // No boundary ≤ LCP → mid-message cut (tools-before-system layout). + TEST_ASSERT(PinFriendlyPrompt::choose_pin_end(80, boundaries, 50) == 80); + TEST_ASSERT(PinFriendlyPrompt::choose_pin_end(40, boundaries, 50) == 0); +} + +TEST_CASE(ServerUnitFixture, test_ppp_annotate_against_recent_ring) { + // Shared tools+identity head, divergent session clock in the tail of the + // first turn (before any chat boundary at 200). + std::vector day1(180, 7); + day1.push_back(111); // date token + day1.insert(day1.end(), {8, 8, 8}); // past boundary material + std::vector day2(180, 7); + day2.push_back(222); + day2.insert(day2.end(), {8, 8, 8}); + + std::vector> ring = {day1}; + const std::vector boundaries = {200}; + const int pin = PinFriendlyPrompt::annotate_pin_end( + day2, boundaries, ring, /*window=*/4, /*min=*/50); + TEST_ASSERT(pin == 180); // mid-message LCP before date drift +} + +TEST_CASE(ServerUnitFixture, test_ppp_diff_split_finds_middle_hunk) { + const std::vector a = {1, 2, 3, 100, 4, 5}; + const std::vector b = {1, 2, 3, 999, 4, 5}; + const auto split = PinFriendlyPrompt::diff_split(a, b); + TEST_ASSERT(split.prefix_len == 3); + TEST_ASSERT(split.suffix_len == 2); + TEST_ASSERT(split.middle_begin == 3); + TEST_ASSERT(split.middle_end == 4); +} + +TEST_CASE(ServerUnitFixture, test_ppp_diff_rewrite_moves_volatile_after_stable) { + // Head: [stable…][TIME][stable_tail…][im_end] → [stable…][stable_tail…][TIME][im_end] + std::vector day1 = {7, 7, 7, 7, 111, 8, 8, 50}; // 50 = im_end + std::vector day2 = {7, 7, 7, 7, 222, 8, 8, 50}; + // Transcript after first boundary. + day2.insert(day2.end(), {9, 9}); + + ChatMarkers markers; + markers.family = "test"; + markers.end_msg_seqs = {{50}}; + + std::vector> ring = {day1}; + // Chat boundaries sit after the next role-start; DiffPin must still cut + // the rewrite head at the first im_end (index 8), not at boundaries.front(). + const std::vector boundaries = {10}; + auto rw = PinFriendlyPrompt::diff_make_pin_friendly( + day2, boundaries, ring, markers, + /*window=*/4, /*min_pin=*/4, /*max_ephemeral=*/16); + TEST_ASSERT(rw.rewritten); + TEST_ASSERT(rw.prefix_len == 4); + TEST_ASSERT(rw.suffix_len == 2); // {8,8} after peeling im_end trailer + TEST_ASSERT(rw.middle_len == 1); + // pin covers stable prefix+suffix; volatile then im_end follow. + TEST_ASSERT(rw.pin_end == 6); + // [7,7,7,7][8,8][222][50][9,9] + TEST_ASSERT(rw.tokens.size() == day2.size()); + TEST_ASSERT((rw.tokens[0] == 7 && rw.tokens[3] == 7)); + TEST_ASSERT(rw.tokens[4] == 8 && rw.tokens[5] == 8); + TEST_ASSERT(rw.tokens[6] == 222); + TEST_ASSERT(rw.tokens[7] == 50); + TEST_ASSERT(rw.tokens[8] == 9); +} + +TEST_CASE(ServerUnitFixture, test_ppp_diff_rewrite_noop_without_boundaries) { + std::vector day1 = {7, 7, 7, 7, 111, 8, 8, 50}; + std::vector day2 = {7, 7, 7, 7, 222, 8, 8, 50, 9, 9}; + ChatMarkers markers; + markers.end_msg_seqs = {{50}}; + std::vector> ring = {day1}; + auto rw = PinFriendlyPrompt::diff_make_pin_friendly( + day2, /*boundaries=*/{}, ring, markers, + /*window=*/4, /*min_pin=*/4, /*max_ephemeral=*/16); + TEST_ASSERT(!rw.rewritten); + TEST_ASSERT(rw.tokens == day2); +} + +TEST_CASE(ServerUnitFixture, test_ppp_diff_rewrite_stops_before_next_role) { + // Realistic boundary: after user role-start (token 90), past im_end (50). + // Volatile middle must not float into the user turn. + std::vector day1 = {7, 7, 7, 7, 111, 8, 8, 50}; + std::vector day2 = {7, 7, 7, 7, 222, 8, 8, 50, 90, 91, 92}; + ChatMarkers markers; + markers.end_msg_seqs = {{50}}; + std::vector> ring = {day1}; + const std::vector boundaries = {11}; // after user role start + auto rw = PinFriendlyPrompt::diff_make_pin_friendly( + day2, boundaries, ring, markers, + /*window=*/4, /*min_pin=*/4, /*max_ephemeral=*/16); + TEST_ASSERT(rw.rewritten); + TEST_ASSERT(rw.tokens.size() == day2.size()); + // Head rewritten; user role tokens untouched at the end. + TEST_ASSERT(rw.tokens[rw.tokens.size() - 3] == 90); + TEST_ASSERT(rw.tokens[rw.tokens.size() - 2] == 91); + TEST_ASSERT(rw.tokens[rw.tokens.size() - 1] == 92); + TEST_ASSERT(rw.tokens[6] == 222); + TEST_ASSERT(rw.tokens[7] == 50); +} + +TEST_CASE(ServerUnitFixture, test_ppp_tools_system_head_end) { + ChatMarkers markers; + markers.end_msg_seqs = {{50}, {51, 52}}; + std::vector ids = {1, 2, 50, 90, 91}; + TEST_ASSERT(PinFriendlyPrompt::tools_system_head_end(ids, markers) == 3); + ids = {1, 2, 51, 52, 90}; + TEST_ASSERT(PinFriendlyPrompt::tools_system_head_end(ids, markers) == 4); + TEST_ASSERT(PinFriendlyPrompt::tools_system_head_end({1, 2, 3}, markers) == 0); +} + +TEST_CASE(ServerUnitFixture, test_ppp_split_and_rearrange_ephemeral_tail) { + const std::string system = + "You are Hermes.\n\n" + "Conversation started: Thursday, July 30, 2026 03:59 PM\n" + "Model: qwen\n"; + auto [stable, ephemeral] = + PinFriendlyPrompt::split_ephemeral_system_tail(system); + TEST_ASSERT(stable == "You are Hermes."); + TEST_ASSERT(ephemeral.find("Conversation started:") == 0); + + std::vector messages = { + {"system", system, ""}, + {"user", "hi", ""}, + }; + auto off = PinFriendlyPrompt::rearrange(messages, false); + TEST_ASSERT(!off.rearranged); + TEST_ASSERT(off.messages.size() == 2); + + auto on = PinFriendlyPrompt::rearrange(messages, true); + TEST_ASSERT(on.rearranged); + TEST_ASSERT(on.messages.size() == 3); + TEST_ASSERT(on.messages[0].role == "system"); + TEST_ASSERT(on.messages[0].content == "You are Hermes."); + TEST_ASSERT(on.messages[1].role == "system"); + TEST_ASSERT(on.messages[1].content.find("Conversation started:") == 0); + TEST_ASSERT(on.messages[2].role == "user"); +} + // ── Prefix-aware eviction policy (model-free) ─────────────────────────── TEST_CASE(ServerUnitFixture, test_evict_empty_is_zero) { @@ -1833,6 +2011,19 @@ TEST_CASE(ServerUnitFixture, test_evict_branch_spares_shared_root) { TEST_ASSERT(v != 0); // the shared root must be spared } +TEST_CASE(ServerUnitFixture, test_evict_skips_protected_leaf) { + // Two unrelated leaves; oldest is protected → evict next unprotected leaf. + std::vector> ids = {{1, 1}, {2, 2}, {3, 3}}; + std::vector protect = {true, false, false}; + TEST_ASSERT(select_inline_evict_victim(ids, &protect) == 1); +} + +TEST_CASE(ServerUnitFixture, test_evict_all_protected_falls_back) { + std::vector> ids = {{1, 1}, {2, 2}}; + std::vector protect = {true, true}; + TEST_ASSERT(select_inline_evict_victim(ids, &protect) == 0); +} + // ═══════════════════════════════════════════════════════════════════════ // PFlash config tests (model-free) // ═══════════════════════════════════════════════════════════════════════