From 86e5837f1207f0de49a0161432cb928ba19528e1 Mon Sep 17 00:00:00 2001
From: Afloat <ShuffleMonkey@163.com>
Date: Wed, 11 Nov 2026 17:39:34 +0800
Subject: [PATCH] perf(text): skip forbidden control-token scans for plain
 input

---
 edge0-infer/src/app/text/bpe.cpp          |  7 ++++++-
 edge0-infer/tests/unit/app/text_tests.cpp | 15 +++++++++++++++
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/edge0-infer/src/app/text/bpe.cpp b/edge0-infer/src/app/text/bpe.cpp
index 724f46cc..8cb66baf 100644
--- a/edge0-infer/src/app/text/bpe.cpp
+++ b/edge0-infer/src/app/text/bpe.cpp
@@ -157,10 +157,12 @@ std::vector<int32_t>
 Bpe::encode_fragments(const std::vector<Fragment> &fragments,
                       bool standalone) const {
   std::string text;
+  const bool has_protocol = std::any_of(fragments.begin(), fragments.end(),
+      [](const Fragment& f) { return f.protocol && !f.text.empty(); });
   std::vector<bool> allowed;
   for (const auto &f : fragments) {
     text += f.text;
-    allowed.insert(allowed.end(), f.text.size(), f.protocol);
+    if (has_protocol) allowed.insert(allowed.end(), f.text.size(), f.protocol);
   }
   if (!valid_utf8(text))
     throw TextError(Error::invalid_argument, "input is not valid UTF-8");
@@ -170,6 +172,9 @@ Bpe::encode_fragments(const std::vector<Fragment> &fragments,
   };
   std::vector<Span> spans = {{0, text.size(), -1}};
   for (auto id : specials_) {
+    // Plain user input cannot contain control tokens. User-defined tokens
+    // still match in standalone encoding, as required by that API.
+    if (!has_protocol && !(standalone && data_.types[id] == 4)) continue;
     const auto &s = data_.tokens[id];
     std::vector<Span> next;
     for (auto span : spans) {
diff --git a/edge0-infer/tests/unit/app/text_tests.cpp b/edge0-infer/tests/unit/app/text_tests.cpp
index 55916dd5..010acea2 100644
--- a/edge0-infer/tests/unit/app/text_tests.cpp
+++ b/edge0-infer/tests/unit/app/text_tests.cpp
@@ -29,6 +29,21 @@ void bpe_test() {
     }
     require(rejected,"unknown pre-tokenizer reports actual and supported names");
   }
+  {
+    auto many=d;
+    for (int i=0;i<512;++i) {
+      many.tokens.push_back("<control_"+std::to_string(i)+">");many.types.push_back(3);
+    }
+    Bpe controls(many);
+    const std::string body="<control_7> custom <control_511>";
+    require(controls.encode_input({{body,false},{"",true}})==Bpe(d).encode_input({{body,false}}),
+            "untrusted control spellings remain ordinary across an empty protocol fragment");
+    const auto standalone=controls.encode(body,false);
+    require(std::find(standalone.begin(),standalone.end(),262)!=standalone.end(),
+            "standalone user-defined tokens remain recognized without protocol input");
+    require(controls.encode("<control_7>",true)==std::vector<int32_t>({270}),
+            "explicit protocol parsing remains enabled");
+  }
   Bpe b(d);
   require(b.encode("abc", false) == std::vector<int32_t>({97, 257}),
           "rank priority");