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

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
@@ -430,9 +454,11 @@ def handle_generate(req_id, params):
             **kwargs,
         )
 
+    should_stop = False
     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 +477,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,22 +486,41 @@ 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,
         "interrupted": False,
     }
     if last_response:
+        generation_tokens = getattr(last_response, "generation_tokens", 0)
+        finish_reason = getattr(last_response, "finish_reason", None)
         result_data.update({
             "prompt_tokens": getattr(last_response, "prompt_tokens", 0),
             "prompt_tps": getattr(last_response, "prompt_tps", 0),
-            "generation_tokens": getattr(last_response, "generation_tokens", 0),
+            "generation_tokens": generation_tokens,
             "generation_tps": getattr(last_response, "generation_tps", 0),
             "peak_memory": getattr(last_response, "peak_memory", 0),
-            "finish_reason": getattr(last_response, "finish_reason", None),
+            "finish_reason": finish_reason,
+            # Same name as llama.node's parallel result, which the completion routes read.
+            # mlx_lm reports "length" at max_tokens; mlx_vlm reports no reason, so a run
+            # that produced max_tokens tokens without a stop sequence hit the limit.
+            "stopped_limit": not should_stop and (
+                finish_reason == "length"
+                or (finish_reason is None and 0 < max_tokens <= generation_tokens)
+            ),
         })
 
     send({"id": req_id, "event": "result", "data": result_data})