serve: tokenise the prompt from the last shared prefix, not from the start

SECURITY.md

@@ -39,6 +39,8 @@ this PC only. The details and every setting are in [docs/DETAILS.md](docs/DETAIL
 - **MCP tools** are opt-in: only the servers you put in the run config, only for requests from Strata's own page
   that ask for them. They run with your user's rights, and the model decides when to call them.
 - **The request monitor** (`/api-monitor`, which keeps the last prompts and answers in memory) is off unless
-  `"api_monitor": true` is set.
+  `"api_monitor": true` is set. Apart from it, the server keeps in its process memory the text and tokens of its
+  last few prompts, whatever `"api_monitor"` says, to encode the next turn faster; no endpoint shows them, and they
+  are gone when the server stops.
 
 Strata has not had an outside security audit yet.

serve/server.py

@@ -551,6 +551,19 @@ def poll(self, path, start) -> dict:
         return dict(self.state)
 
 
+def conversation_slots(engine) -> int | None:
+    """How many conversations the engine's cache holds: its --conversation-cache-slots argument, which is known
+    before the engine has started (a lazy start has reported no INFO line yet), else the INFO line it reported."""
+    args = engine.spawn[1] if hasattr(engine, "spawn") else []
+    if "--conversation-cache-slots" in args[:-1]:
+        try:
+            return int(args[args.index("--conversation-cache-slots") + 1])
+        except ValueError:
+            pass
+    slots = (getattr(engine, "info", None) or {}).get("conversation_cache_slots")
+    return slots if isinstance(slots, int) else None
+
+
 def conversation_cache_view(info: dict, hist: list, totals: dict, parked: dict) -> dict:
     """#596: the Monitor's Conversation cache card: the parked conversations (the engine's opt-in
     --conversation-cache-mib: budget, slots, what its log says) and how much of the prompts the cache gave back."""
@@ -2567,23 +2580,32 @@ class ByteTokenizer:
 
     ALWAYS = ()                                     # specials matched without parse_special (type 4, as <think>)
 
+    @property
+    def max_special_len(self):
+        return max(len(s) for s in self.SPECIALS)
+
     @property
     def control_tokens(self):
         return [s for s in self.SPECIALS if s not in self.ALWAYS]
 
     def encode(self, text, parse_special=False, plain=()):
-        out, i = [], 0
+        return self.encode_marked(text, parse_special, plain)[0]
+
+    def encode_marked(self, text, parse_special=False, plain=()):
+        """encode() and its resume points (after each special), as strata_tokenizer's for PromptEncoder."""
+        out, marks, i = [], [], 0
         while i < len(text):
             for k, s in enumerate(self.SPECIALS):
                 if (parse_special or s in self.ALWAYS) and text.startswith(s, i) and not any(
                         a <= i < b for a, b in plain):
                     out.append(256 + k)
                     i += len(s)
+                    marks.append((i, len(out)))
                     break
             else:
                 out.extend(text[i].encode("utf-8"))
                 i += 1
-        return out
+        return out, marks
 
     def decode(self, ids, errors="replace"):
         raw = bytearray()
@@ -2794,6 +2816,17 @@ def __init__(self, engine: Engine, tokenizer, template: ChatTemplate, model_name
         self.tool_call_recovery = False                   # opt-in: tool calls in forms next to the template's
         self.stop_ids = set(tokenizer.encode(IM_END, parse_special=True) +
                             tokenizer.encode("<|endoftext|>", parse_special=True))
+        # #567: a prompt re-encodes only what follows the last special token it shares with a recent prompt (the
+        # same ids as a full encode: tools/strata_tokenizer.py PromptEncoder).  Tokenizers without resume points
+        # encode in full.
+        self.prompts = None
+        if hasattr(tokenizer, "encode_marked") and hasattr(tokenizer, "max_special_len"):
+            from strata_tokenizer import PromptEncoder
+            # one kept prompt per conversation the engine's cache can hold (--conversation-cache-slots); the
+            # encoder's own default without that setting
+            slots = conversation_slots(engine)
+            self.prompts = PromptEncoder(tokenizer, keep=slots) if slots is not None and slots > 0 \
+                else PromptEncoder(tokenizer)
 
     def loaded(self) -> bool:
         return not hasattr(self.engine, "alive") or self.engine.alive()
@@ -3349,17 +3382,21 @@ def render_prompt(self, messages, tools, kwargs) -> str:
         tail = "" if off else EFFORT_TURN.format(EFFORT_TEXT.get(effort, EFFORT_TEXT["medium"]))
         return history + tail + full[len(mine):]
 
-    def encode_prompt(self, messages, tools, kwargs) -> list[int]:
+    def encode_prompt(self, messages, tools, kwargs, store: bool = True) -> list[int]:
         """The request's prompt: the template rendered and tokenized.  #537: a <think> / </think> written inside a
         message's text is encoded as the text it is, not as the model's reasoning markers (the template's own are).
         A control token's text (<|im_start|>, <|im_end|>, <|endoftext|>, ...) inside a message is text as well: only
-        the control tokens the template writes are control tokens."""
+        the control tokens the template writes are control tokens.  `store=False` (a token count) reuses the kept
+        prompts without adding or replacing one."""
         marked, marked_tools, changed = mark_think_literals(messages, tools, self.literals)
         prompt = self.render_prompt(marked, marked_tools, kwargs)
-        if not changed:
-            return self.tok.encode(prompt, parse_special=True)
-        prompt, plain = unmark_think_literals(prompt, self.literals)
-        return self.tok.encode(prompt, parse_special=True, plain=plain)
+        plain = ()
+        if changed:
+            prompt, plain = unmark_think_literals(prompt, self.literals)
+        if self.prompts is not None:            # #567: only what follows the last special shared with a recent prompt
+            return self.prompts.encode(prompt, plain, store)
+        return self.tok.encode(prompt, parse_special=True, plain=plain) if plain else \
+            self.tok.encode(prompt, parse_special=True)
 
     def _note_unreadable_tool_images(self, messages):
         """A picture a tool returned (Claude Code's Read of an image file) that this server cannot read - it has no
@@ -3390,12 +3427,16 @@ def _note_unreadable_tool_images(self, messages):
                     content[n] = {"type": "text", "text": f"[image omitted: {why}]"}
         return fetched
 
+    def forget_prompt(self, ids):
+        if self.prompts is not None:
+            self.prompts.forget(ids)
+
     def prepare(self, messages, tools, kwargs, max_new=None, force=None, req=None):
         """-> (ids, thinking, max_new). An unset or non-positive max_new (some clients send -1) means "unlimited":
         the rest of the context.  `force` (forced_call): without thinking the reply starts with it, so it ends the
         prompt; with thinking, Service.run writes it once the thinking is over."""
         fetched = self._note_unreadable_tool_images(messages)
-        ids = self.encode_prompt(messages, tools, kwargs)
+        ids = encoded_ids = self.encode_prompt(messages, tools, kwargs)
         if force and kwargs.get("enable_thinking", True) is False:
             ids = ids + self.tok.encode(force, parse_special=True)
         self.embeddings.path = None
@@ -3448,12 +3489,16 @@ def prepare(self, messages, tools, kwargs, max_new=None, force=None, req=None):
         room = ctx - CTX_SLACK - len(ids)
         if max_new is None or max_new <= 0 or (self.fit_max_tokens and room < 1):
             if room < 1:
+                # A prompt refused for the context is not kept (#567): the client cannot resend it unchanged.  The
+                # other refusals (engine starting, images) keep it: the same prompt comes back and its encode is reused.
+                self.forget_prompt(encoded_ids)
                 raise ValueError("request exceeds the context window: "
                                  f"prompt ({len(ids)} tokens) leaves no room to answer in the context "
                                  f"({ctx}); requests are never truncated")
             max_new = room
         elif max_new > room:
             if not self.fit_max_tokens:
+                self.forget_prompt(encoded_ids)
                 raise ValueError("request exceeds the context window: "
                                  f"prompt ({len(ids)} tokens) + max tokens ({max_new}) exceeds the context "
                                  f"({ctx}); requests are never truncated. Send a smaller "
@@ -5446,7 +5491,7 @@ def _count_tokens(self, req):
             messages, tools, kw = anthropic_to_messages(req, svc.anthropic_think_unasked)
             stop_strings(req)                                 # the same 400 as the request itself would get
             check_request_sampling(req)                       # ... likewise a sampling field of the wrong type
-            self._json(200, {"input_tokens": len(svc.encode_prompt(messages, tools, kw))})
+            self._json(200, {"input_tokens": len(svc.encode_prompt(messages, tools, kw, store=False))})
 
         def _anthropic(self, req):
             svc.load()

serve/test_prompt_encoder.py

@@ -0,0 +1,162 @@
+"""serve/test_prompt_encoder.py - #567's incremental prompt encoding with #537's plain spans: every prompt's ids
+are those of a full encode, and a turn reuses the previous one's.
+
+    python -m unittest serve.test_prompt_encoder -v
+"""
+from __future__ import annotations
+
+import random
+import re
+import sys
+import unittest
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+sys.path.insert(0, str(ROOT / "tools"))
+from serve.frontend import ChatTemplate  # noqa: E402
+from serve.server import ByteTokenizer, MockEngine, Service  # noqa: E402
+from strata_tokenizer import PromptEncoder, plain_agree_until  # noqa: E402
+
+
+class ThinkTokenizer(ByteTokenizer):
+    """ByteTokenizer whose <think> / </think> are specials matched even without parse_special (type 4), as in the
+    Qwen vocabulary: the case #537's plain spans exist for."""
+    SPECIALS = ByteTokenizer.SPECIALS + ["<think>", "</think>"]
+    ALWAYS = ("<think>", "</think>")
+
+
+PIECES = ["<|im_start|>user\n", "<|im_end|>\n", "<think>", "</think>", "hello ", "é你 ", "<|im_sta", "x" * 7,
+          "<|im_start|>assistant\n", "</thi", "\n\n", "<|endoftext|>", "<|vision_start|>", "<|image_pad|>"]
+
+
+def spans_of(text, rng):
+    """Random plain spans over the specials' texts (<think>, </think>, <|im_start|>, ...), as unmark_think_literals
+    marks the ones quoted inside a message."""
+    matches = re.finditer("|".join(map(re.escape, ThinkTokenizer.SPECIALS)), text)
+    return tuple((m.start(), m.end()) for m in matches if rng.random() < 0.5)   # finditer: left to right, no overlap
+
+
+class Units(unittest.TestCase):
+    def test_plain_agree_until(self):
+        self.assertEqual(plain_agree_until((), (), 50), 50)
+        self.assertEqual(plain_agree_until(((5, 13),), ((5, 13),), 50), 50)
+        self.assertEqual(plain_agree_until(((5, 13),), (), 50), 5)
+        self.assertEqual(plain_agree_until(((5, 13),), ((5, 9),), 50), 9)
+        self.assertEqual(plain_agree_until(((60, 70),), (), 50), 50)          # beyond the shared text: no matter
+        self.assertEqual(plain_agree_until(((5, 13), (20, 28)), ((5, 13),), 50), 20)
+
+    def test_same_ids_as_a_full_encode(self):
+        tok = ThinkTokenizer()
+        rng = random.Random(7)
+        for conv in range(300):
+            enc = PromptEncoder(tok)
+            text, plain = "", ()
+            for turn in range(8):
+                text += "".join(rng.choice(PIECES) for _ in range(rng.randint(1, 12)))
+                if rng.random() < 0.6:
+                    plain = spans_of(text, rng)                     # the spans may change in the shared part too
+                with self.subTest(conv=conv, turn=turn):
+                    self.assertEqual(enc.encode(text, plain), tok.encode(text, True, plain))
+                if rng.random() < 0.2 and len(text) > 10:              # an edit: a client re-rendering its history
+                    cut = rng.randrange(len(text))
+                    text = text[:cut]
+                    plain = tuple((a, b) for a, b in plain if b <= cut)
+
+    def test_a_turn_reuses_the_previous_one(self):
+        tok = ThinkTokenizer()
+        enc = PromptEncoder(tok)
+        # as a chat turn: the earlier prompt goes on past its last special (the answer being written), and the
+        # next one shares it up to there - the resume point needs max_special_len characters after it
+        base = "<|im_start|>user\n" + "quoted </think> here " * 200 + "<|im_end|>\n<|im_start|>assistant\nan answer"
+        plain = tuple((i, i + 8) for i in range(len(base)) if base.startswith("</think>", i))
+        enc.encode(base, plain)
+        more = base + " that goes on<|im_end|>\n<|im_start|>user\nnext<|im_end|>\n"
+        self.assertEqual(enc.encode(more, plain), tok.encode(more, True, plain))
+        self.assertGreater(enc.last_reused, len(base) // 2)
+        # the same text with the quoted </think> no longer plain: nothing before it may be reused past its start
+        self.assertEqual(enc.encode(more, ()), tok.encode(more, True, ()))
+        self.assertLessEqual(enc.last_reused, plain[0][0])
+
+
+class Served(unittest.TestCase):
+    def test_quoted_think_across_turns(self):
+        tok = ThinkTokenizer()
+        template = ChatTemplate(ROOT / "serve/chat_template.jinja")
+        svc = Service(MockEngine(tok, "</think>\n\nok", max_context=1 << 20), tok, template)
+        full = Service(MockEngine(tok, "</think>\n\nok", max_context=1 << 20), tok, template)
+        full.prompts = None                                        # the reference: a full encode every time
+        self.assertIsNotNone(svc.prompts)
+        msgs = [{"role": "system", "content": "Explain tags."}]
+        for turn in range(5):
+            msgs.append({"role": "user", "content": f"turn {turn}: what does </think> mean? <think> too " * 3})
+            got, want = svc.encode_prompt(msgs, None, {}), full.encode_prompt(msgs, None, {})
+            self.assertEqual(got, want)
+            if turn:
+                self.assertGreater(svc.prompts.last_reused, 0)
+            msgs.append({"role": "assistant", "content": "It closes the reasoning: </think> is a marker."})
+
+    def service(self, max_context=1 << 22):
+        tok = ByteTokenizer()
+        return Service(MockEngine(tok, "ok", max_context=max_context), tok, ChatTemplate(ROOT / "serve/chat_template.jinja"))
+
+    def test_another_conversation_does_not_evict_the_first(self):
+        """A subagent beside the main conversation: every rendered prompt starts with the template's header, which
+        must not make the second client's prompt a continuation of the first's entry."""
+        svc = self.service()
+        turn2 = [{"role": "user", "content": "a long question " * 500}, {"role": "assistant", "content": "answer"},
+                 {"role": "user", "content": "next"}]
+        svc.encode_prompt(turn2, None, {})
+        svc.encode_prompt([{"role": "user", "content": "one unrelated prompt " * 300}], None, {})
+        self.assertEqual(len(svc.prompts.entries), 2)
+        turn3 = turn2 + [{"role": "assistant", "content": "x"}, {"role": "user", "content": "more"}]
+        svc.encode_prompt(turn3, None, {})
+        self.assertGreater(svc.prompts.last_reused, len(turn2[0]["content"]))
+        self.assertEqual(len(svc.prompts.entries), 2)           # turn 3 replaced turn 2, the other one stayed
+
+    def test_a_refused_prompt_is_not_kept(self):
+        svc = self.service(max_context=4096)
+        msgs = [{"role": "user", "content": "kept " * 20}]
+        svc.prepare(msgs, None, {}, 16)
+        refused = msgs + [{"role": "assistant", "content": "a"}, {"role": "user", "content": "more " * 20}]
+        full = self.service(max_context=4096)
+        full.prompts = None
+        refused_ids = full.encode_prompt(refused, None, {})
+        with self.assertRaisesRegex(ValueError, "exceeds the context"):
+            svc.prepare(refused, None, {}, 1 << 20)
+        self.assertNotIn(refused_ids, [e[2] for e in svc.prompts.entries])
+        self.assertNotIn(full.render_prompt(refused, None, {}), [e[0] for e in svc.prompts.entries])
+
+    def test_random_messages_equal_a_full_encode(self):
+        """Messages quoting control-token texts and <think> tags at random, growing, edited and cut back between
+        requests: the incremental encoder's ids are the full encode's, 2400 prompts."""
+        tok = ThinkTokenizer()
+        template = ChatTemplate(ROOT / "serve/chat_template.jinja")
+        svc = Service(MockEngine(tok, "ok", max_context=1 << 20), tok, template)
+        full = Service(MockEngine(tok, "ok", max_context=1 << 20), tok, template)
+        full.prompts = None
+        rng = random.Random(567)
+        quoted = ThinkTokenizer.SPECIALS + ["<|im_sta", "</thi"]
+        words = ["hello ", "é你 ", "x" * 9, "\n", "tool "]
+
+        def text():
+            return "".join(rng.choice(quoted if rng.random() < 0.4 else words) for _ in range(rng.randint(1, 10)))
+
+        reused = 0
+        for conv in range(300):
+            msgs = [{"role": "system", "content": text()}]
+            for turn in range(8):
+                msgs.append({"role": "user", "content": text()})
+                if rng.random() < 0.15:                             # an earlier turn edited by the client
+                    msgs[rng.randrange(len(msgs))]["content"] = text()
+                if rng.random() < 0.1 and len(msgs) > 2:            # a shorter conversation
+                    msgs = msgs[:rng.randrange(1, len(msgs))] + msgs[-1:]
+                with self.subTest(conv=conv, turn=turn):
+                    self.assertEqual(svc.encode_prompt(msgs, None, {}), full.encode_prompt(msgs, None, {}))
+                reused += svc.prompts.last_reused > 0
+                msgs.append({"role": "assistant", "content": text()})
+        self.assertGreater(reused, 1000)
+
+
+if __name__ == "__main__":
+    unittest.main()