From 482ecff229258fa9635c36eb2cbcdf283e8015f9 Mon Sep 17 00:00:00 2001
From: Chris Greening <cgreening2@apple.com>
Date: Wed, 19 Aug 2026 15:59:26 +0100
Subject: [PATCH] Fix guided generation: don't leak grammar terminal token on
 GPU path

The GPU constrained decode loop (runConstrainedCompletion) yielded each
sampled token to the consumer *before* asking the grammar whether that
token terminated it. When the grammar terminates on a stop token (e.g.
<|endoftext|>, 151643), that token had already been streamed and got
decoded into the structured output, producing trailing text like:

    {...}<|endoftext|>

This is the GPU analogue of the CPU bug fixed in #117, which only patched
ConstrainedDecodingStrategy. The GPU path added later in #170 reintroduced
the same class of bug and was never covered by #117's fix.

Defer the yield: a sampled token is now emitted only after the grammar has
accepted it and we've confirmed the acceptance did not terminate. This
matches the accept -> isTerminated -> yield ordering already used by the
test mock (MockConstrainedEngine).

Verified with llm-runner --json-schema against Qwen3-0.6B (greedy):
  before: {...}<|endoftext|>
  after:  {...}
---
 .../CoreAIPipelinedEngine.swift               | 29 ++++++++++++++-----
 1 file changed, 21 insertions(+), 8 deletions(-)

diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift
index 000a4416..6c8f78a3 100644
--- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift
+++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift
@@ -1247,18 +1247,31 @@ private struct EngineImpl: ~Copyable {
             }
         }
 
-        continuation.yield(lastToken)
+        // NOTE: `lastToken` is deliberately NOT yielded here. A sampled token is
+        // only emitted after the grammar has accepted it (at the top of the loop)
+        // and we've confirmed the acceptance did not terminate the grammar.
+        // Otherwise the grammar's terminal token (e.g. <|endoftext|>, 151643) would
+        // be streamed to the consumer before termination is detected and would leak
+        // into the structured output. This mirrors the CPU ConstrainedDecodingStrategy
+        // fix in #117, which returns nil instead of emitting the terminal token's text.
 
         // Constrained decode loop
-        var generated = 1
+        var generated = 0
         while generated < maxTokens {
             guard !isCancelled.load(ordering: .relaxed) else { break }
             try Task.checkCancellation()
 
-            // Accept previous token in grammar
+            // Accept the most recently sampled token in the grammar.
             if !session.acceptToken(lastToken) { break }
+            // If accepting it terminated the grammar, `lastToken` is the terminal
+            // (stop) token — break WITHOUT emitting it.
             if session.isTerminated { break }
 
+            // `lastToken` is confirmed to be valid structured content — emit it now.
+            continuation.yield(lastToken)
+            generated += 1
+            if generated >= maxTokens { break }
+
             // Check for jump-forward: deterministic grammar segments that can be
             // batch-encoded without per-token sampling (saves N-1 round-trips)
             if let jumpString = session.findJumpForwardString(),
@@ -1288,12 +1301,13 @@ private struct EngineImpl: ~Copyable {
                     }
                 }
 
-                // Yield the jump-forward tokens (deterministic) + the sampled token
+                // Emit the deterministic jump-forward tokens now. The freshly sampled
+                // `lastToken` is deferred to the next iteration's accept/termination
+                // check so a terminal token never leaks into the output.
                 for jt in jumpTokens {
                     continuation.yield(jt)
                 }
-                continuation.yield(lastToken)
-                generated += jumpTokens.count + 1
+                generated += jumpTokens.count
                 continue
             }
 
@@ -1322,8 +1336,7 @@ private struct EngineImpl: ~Copyable {
                 }
             }
 
-            continuation.yield(lastToken)
-            generated += 1
+            // `lastToken` is deferred to the next iteration's accept/termination check.
         }
 
         // Drain: sentinel command buffer to ensure all GPU work completes