@fugood/buttress-server 2.26.0-beta.20 → 2.26.0-beta.21

package/lib/mlx-bridge.py

@@ -242,6 +242,17 @@ def clear_cancel(req_id):
         _cancel_flags.discard(req_id)
 
 
+def _partial_stop_length(text, stop_sequences):
+    """Length of the longest end of `text` that a stop sequence starts with."""
+    longest = 0
+    for seq in filter(None, stop_sequences):
+        for length in range(min(len(seq) - 1, len(text)), longest, -1):
+            if text.endswith(seq[:length]):
+                longest = length
+                break
+    return longest
+
+
 def _get_tokenizer():
     """Return the underlying tokenizer from the processor."""
     if _processor is None:
@@ -407,6 +418,19 @@ def handle_generate(req_id, params):
 
     full_text = ""
     last_response = None
+    # Token events carry full_text[:sent_end]. The client builds its output
+    # from them, so they must never contain stop sequence text.
+    sent_end = 0
+
+    def send_text_until(end, token_id):
+        nonlocal sent_end
+        if end > sent_end:
+            send({
+                "id": req_id,
+                "event": "token",
+                "data": {"token": full_text[sent_end:end], "token_id": token_id},
+            })
+            sent_end = end
 
     if _is_vlm:
         from mlx_vlm import stream_generate
@@ -433,6 +457,7 @@ def handle_generate(req_id, params):
     for response in gen_iter:
         if is_cancelled(req_id):
             clear_cancel(req_id)
+            send_text_until(len(full_text), None)
             send({
                 "id": req_id,
                 "event": "result",
@@ -451,15 +476,6 @@ def handle_generate(req_id, params):
         full_text += response.text
         last_response = response
 
-        send({
-            "id": req_id,
-            "event": "token",
-            "data": {
-                "token": response.text,
-                "token_id": getattr(response, "token", None),
-            },
-        })
-
         # Check stop sequences
         should_stop = False
         for seq in stop_sequences:
@@ -469,9 +485,19 @@ def handle_generate(req_id, params):
                 idx = full_text.index(seq)
                 full_text = full_text[:idx]
                 break
+
+        # Like llama.node, hold back an end that may still become a stop
+        # sequence until a later token completes or breaks it.
+        safe_end = len(full_text)
+        if not should_stop:
+            safe_end -= _partial_stop_length(full_text, stop_sequences)
+        send_text_until(safe_end, getattr(response, "token", None))
         if should_stop:
             break
 
+    # Generation ended without completing the held-back stop sequence
+    send_text_until(len(full_text), None)
+
     # Send final result
     result_data = {
         "text": full_text,