Keep a grammar applied inside a thinking-budget block on request

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- `CompletionOptions.thinkingBudget.pauseGrammar: false` keeps a grammar applied inside a reasoning block, so a grammar that takes the block and its end tag itself, such as one for a prompt that ends inside the block, sees the budget's forced `</think>`. It needs the new `thinkingBudgetPauseGrammar` capability; omitted, the grammar pauses as before.
 - WebGPU text-to-speech applies the GGUF's suppress tokens in the talker sampler, as upstream does.
   Qwen3-TTS no longer fails with `code out of range` when a codec control token is sampled.
 - `createCompletion` can explicitly count prompt tokens in penalties with a grammar using `penaltyPromptTokensWithGrammar`; a matching loaded-core capability is required, and existing grammar defaults stay unchanged.

CMakeLists.txt

@@ -235,7 +235,7 @@ set(LLAMADART_WEBGPU_LINK_OPTIONS
     "-sEXPORT_NAME=createLlamaWebGpuCoreModule"
     "-sENVIRONMENT=web,worker"
     "-sEXPORTED_RUNTIME_METHODS=['FS','ccall','UTF8ToString']"
-    "-sEXPORTED_FUNCTIONS=['_main','_llamadart_webgpu_probe','_llamadart_webgpu_supports_pthreads','_llamadart_webgpu_pthread_pool_size','_llamadart_webgpu_backends_json','_llamadart_webgpu_last_error','_llamadart_webgpu_set_log_level','_llamadart_webgpu_load_model','_llamadart_webgpu_load_model_from_url','_llamadart_webgpu_free_model','_llamadart_webgpu_mmproj_load','_llamadart_webgpu_mmproj_free','_llamadart_webgpu_mmproj_supports_vision','_llamadart_webgpu_mmproj_supports_audio','_llamadart_webgpu_media_clear_pending','_llamadart_webgpu_media_add_file','_llamadart_webgpu_media_add_encoded','_llamadart_webgpu_media_add_rgb','_llamadart_webgpu_media_add_audio_f32','_llamadart_webgpu_tokenize_to_json','_llamadart_webgpu_tokenize_bytes_to_json','_llamadart_webgpu_last_tokens_json','_llamadart_webgpu_state_save_file','_llamadart_webgpu_state_load_file','_llamadart_webgpu_test_state_load_failures','_llamadart_webgpu_detokenize_from_json','_llamadart_webgpu_last_detokenized','_llamadart_webgpu_last_detokenized_bytes_json','_llamadart_webgpu_embed_to_json','_llamadart_webgpu_embed_bytes_to_json','_llamadart_webgpu_embed_parts_to_json','_llamadart_webgpu_embedding_capabilities_json','_llamadart_webgpu_last_embedding_json','_llamadart_webgpu_score_next_token_to_json','_llamadart_webgpu_score_next_token_bytes_to_json','_llamadart_webgpu_last_next_token_scores_json','_llamadart_webgpu_generate','_llamadart_webgpu_begin_generation','_llamadart_webgpu_begin_generation_bytes','_llamadart_webgpu_next_token','_llamadart_webgpu_last_piece','_llamadart_webgpu_end_generation','_llamadart_webgpu_request_cancel','_llamadart_webgpu_last_output','_llamadart_webgpu_last_output_bytes_json','_llamadart_webgpu_last_generation_usage_json','_llamadart_webgpu_completion_capabilities_json','_llamadart_webgpu_set_next_speculative','_llamadart_webgpu_set_next_grammar','_llamadart_webgpu_add_next_grammar_trigger','_llamadart_webgpu_draft_model_load','_llamadart_webgpu_draft_model_free','_llamadart_webgpu_draft_model_info_json','_llamadart_webgpu_heap_headroom_bytes','_llamadart_webgpu_get_context_size','_llamadart_webgpu_model_meta_json','_llamadart_webgpu_tts_api_version','_llamadart_webgpu_tts_info_json','_llamadart_webgpu_tts_start','_llamadart_webgpu_tts_step','_llamadart_webgpu_tts_progress_json','_llamadart_webgpu_tts_write_pcm','_llamadart_webgpu_tts_reset','_llamadart_webgpu_decision_api_version','_llamadart_webgpu_decision_capabilities_json','_llamadart_webgpu_decision_load','_llamadart_webgpu_decision_head_info_json','_llamadart_webgpu_decision_run','_llamadart_webgpu_decision_free','_llamadart_webgpu_lora_api_version','_llamadart_webgpu_lora_load','_llamadart_webgpu_lora_set','_llamadart_webgpu_lora_remove','_llamadart_webgpu_lora_clear','_llamadart_webgpu_shutdown']"
+    "-sEXPORTED_FUNCTIONS=['_main','_llamadart_webgpu_probe','_llamadart_webgpu_supports_pthreads','_llamadart_webgpu_pthread_pool_size','_llamadart_webgpu_backends_json','_llamadart_webgpu_last_error','_llamadart_webgpu_set_log_level','_llamadart_webgpu_load_model','_llamadart_webgpu_load_model_from_url','_llamadart_webgpu_free_model','_llamadart_webgpu_mmproj_load','_llamadart_webgpu_mmproj_free','_llamadart_webgpu_mmproj_supports_vision','_llamadart_webgpu_mmproj_supports_audio','_llamadart_webgpu_media_clear_pending','_llamadart_webgpu_media_add_file','_llamadart_webgpu_media_add_encoded','_llamadart_webgpu_media_add_rgb','_llamadart_webgpu_media_add_audio_f32','_llamadart_webgpu_tokenize_to_json','_llamadart_webgpu_tokenize_bytes_to_json','_llamadart_webgpu_last_tokens_json','_llamadart_webgpu_state_save_file','_llamadart_webgpu_state_load_file','_llamadart_webgpu_test_state_load_failures','_llamadart_webgpu_detokenize_from_json','_llamadart_webgpu_last_detokenized','_llamadart_webgpu_last_detokenized_bytes_json','_llamadart_webgpu_embed_to_json','_llamadart_webgpu_embed_bytes_to_json','_llamadart_webgpu_embed_parts_to_json','_llamadart_webgpu_embedding_capabilities_json','_llamadart_webgpu_last_embedding_json','_llamadart_webgpu_score_next_token_to_json','_llamadart_webgpu_score_next_token_bytes_to_json','_llamadart_webgpu_last_next_token_scores_json','_llamadart_webgpu_generate','_llamadart_webgpu_begin_generation','_llamadart_webgpu_begin_generation_bytes','_llamadart_webgpu_next_token','_llamadart_webgpu_last_piece','_llamadart_webgpu_end_generation','_llamadart_webgpu_request_cancel','_llamadart_webgpu_last_output','_llamadart_webgpu_last_output_bytes_json','_llamadart_webgpu_last_generation_usage_json','_llamadart_webgpu_completion_capabilities_json','_llamadart_webgpu_set_next_speculative','_llamadart_webgpu_set_next_grammar','_llamadart_webgpu_add_next_grammar_trigger','_llamadart_webgpu_set_next_grammar_reasoning_pause','_llamadart_webgpu_draft_model_load','_llamadart_webgpu_draft_model_free','_llamadart_webgpu_draft_model_info_json','_llamadart_webgpu_heap_headroom_bytes','_llamadart_webgpu_get_context_size','_llamadart_webgpu_model_meta_json','_llamadart_webgpu_tts_api_version','_llamadart_webgpu_tts_info_json','_llamadart_webgpu_tts_start','_llamadart_webgpu_tts_step','_llamadart_webgpu_tts_progress_json','_llamadart_webgpu_tts_write_pcm','_llamadart_webgpu_tts_reset','_llamadart_webgpu_decision_api_version','_llamadart_webgpu_decision_capabilities_json','_llamadart_webgpu_decision_load','_llamadart_webgpu_decision_head_info_json','_llamadart_webgpu_decision_run','_llamadart_webgpu_decision_free','_llamadart_webgpu_lora_api_version','_llamadart_webgpu_lora_load','_llamadart_webgpu_lora_set','_llamadart_webgpu_lora_remove','_llamadart_webgpu_lora_clear','_llamadart_webgpu_shutdown']"
     "-lwasmfs_fetch.js"
     "${LLAMADART_WEBGPU_EXCEPTION_CATCHING_ALLOWED}"
     "--post-js=${LLAMADART_WEBGPU_CORE_POST_JS}"

docs/api.md

@@ -326,7 +326,7 @@ Common `options` keys:
 | `grammarLazy` | `true` applies `grammar` only after a `grammarTriggers` entry fires, as llama.cpp's lazy grammar (`llama-server`'s `grammar_lazy`) does; see [Lazy grammars](#lazy-grammars). Needs a non-empty `grammar` and at least one trigger. Omitted or `null` means `false`. |
 | `grammarTriggers` | Array of `{ type: 'word' \| 'pattern' \| 'pattern-full', value }` or `{ type: 'token', token }` that fire a lazy grammar. Always validated; applied only when `grammarLazy` is `true`. |
 | `seed` | Integer seed; random when omitted. |
-| `thinkingBudget` | `{ maxTokens, startTag, endTag, forcedMessage? }`: llama.cpp's reasoning budget. After `maxTokens` tokens inside a block opened by `startTag`, plus any that finish a UTF-8 character, the sampler forces `forcedMessage` (default empty) and `endTag`. A block the prompt leaves open counts from the start of generation, so `maxTokens: 0` forces the end at once. Each new `startTag` gets a fresh budget. A grammar constrains `startTag` like any other text, then pauses until `endTag` completes. `maxTokens` is an integer from `0` to `2147483647`; the tags must be non-empty. Each tag is tokenized on its own, as native llamadart does: with a vocabulary that adds a leading space to text, such as Llama 2's SentencePiece, a tag in the prompt matches only after a space, and the forced end tag starts with one. Text-only prompts: with `parts` it rejects. |
+| `thinkingBudget` | `{ maxTokens, startTag, endTag, forcedMessage?, pauseGrammar? }`: llama.cpp's reasoning budget. After `maxTokens` tokens inside a block opened by `startTag`, plus any that finish a UTF-8 character, the sampler forces `forcedMessage` (default empty) and `endTag`. A block the prompt leaves open counts from the start of generation, so `maxTokens: 0` forces the end at once. Each new `startTag` gets a fresh budget. A grammar constrains `startTag` like any other text, then pauses until `endTag` completes, unless `pauseGrammar` is `false`: that keeps the grammar applied inside the block, for one that takes the block and `endTag` itself (such as one written for a prompt that ends inside the block; paused, it resumes at its start and takes the reply as reasoning) and so accepts `forcedMessage` and `endTag` wherever the budget can run out. Omitted, `null` or `true` pauses it; `false` needs the `thinkingBudgetPauseGrammar` capability. `maxTokens` is an integer from `0` to `2147483647`; the tags must be non-empty. Each tag is tokenized on its own, as native llamadart does: with a vocabulary that adds a leading space to text, such as Llama 2's SentencePiece, a tag in the prompt matches only after a space, and the forced end tag starts with one. Text-only prompts: with `parts` it rejects. |
 | `onUsage(usage)` | Called at most once, with the `CompletionUsage` of the last attempt whose generation returned, before the promise resolves or rejects with an `AbortError`. Not called when no generation returned: other rejections, a skipped multimodal warmup, or a worker failure after a cancel. |
 | `onToken(piece, currentText)` | Token callback. By default `piece` is a `Uint8Array` containing complete UTF-8 characters; see [Streamed text](#streamed-text). `currentText` is the text streamed so far. Direct runtime mode provides it by default; worker mode provides `''` unless `emitCurrentTextOnToken: true` is set. |
 | `signal` | `AbortSignal`; aborting cancels this operation and rejects with `AbortError`. |
@@ -365,7 +365,7 @@ An invalid `minP`, a `presencePenalty` that is not finite as a 32-bit float,
 an invalid `thinkingBudget`, or a malformed `grammarLazy`, `grammarTriggers` or
 `grammarRoot` rejects with a `RangeError` or `TypeError` before generation
 starts, and `grammarLazy: true` or another `grammarRoot` without a non-empty
-`grammar`, or a lazy grammar without triggers, with an `Error`. A nonzero `minP` or `presencePenalty`, any `thinkingBudget`, and
+`grammar`, or a lazy grammar without triggers, with an `Error`. A nonzero `minP` or `presencePenalty`, any `thinkingBudget`, `thinkingBudget.pauseGrammar: false`, and
 `grammarLazy: true` or a `grammarRoot` other than `root` reject when the loaded
 core does not apply that option.
 
@@ -400,7 +400,7 @@ completion rejects with `The lazy grammar rejected the text that triggered it`.
 The core stays usable after either, and worker mode rethrows both without
 falling back to the main thread; like an invalid grammar, each rejection leaks
 a few KiB at most. Triggers match only generated text, not the
-prompt, and a `thinkingBudget` block pauses a lazy grammar, so text inside the
+prompt, and a `thinkingBudget` block pauses a lazy grammar (unless `pauseGrammar` is `false`), so text inside the
 block fires no trigger. The bridge streams every special token's text, so
 llama.cpp's `preserved_tokens` has no equivalent here.
 
@@ -430,7 +430,7 @@ cancelled generation once it returns.
 getCompletionCapabilities(): Promise<CompletionCapabilities>
 ```
 
-Returns `{ presencePenalty, minP, thinkingBudget, penaltyPromptTokens, penaltyPromptTokensWithGrammar, lazyGrammar, speculativeDecoding }`:
+Returns `{ presencePenalty, minP, thinkingBudget, thinkingBudgetPauseGrammar, penaltyPromptTokens, penaltyPromptTokensWithGrammar, lazyGrammar, speculativeDecoding }`:
 whether the loaded core applies each `createCompletion` option. `lazyGrammar`
 covers `grammarLazy`, `grammarTriggers` and `grammarRoot`.
 `penaltyPromptTokensWithGrammar` separately permits the explicit grammar opt-in;

js/llama_webgpu_bridge.d.ts

@@ -241,7 +241,7 @@ export type GrammarTrigger =
   | { type: 'word' | 'pattern' | 'pattern-full'; value: string }
   | { type: 'token'; token: number };
 
-/** llama.cpp's reasoning budget: after `maxTokens` tokens inside a reasoning block, plus any that finish a UTF-8 character, the sampler forces `forcedMessage` and `endTag`. A grammar pauses after `startTag` until `endTag` completes. */
+/** llama.cpp's reasoning budget: after `maxTokens` tokens inside a reasoning block, plus any that finish a UTF-8 character, the sampler forces `forcedMessage` and `endTag`. A grammar pauses after `startTag` until `endTag` completes, unless `pauseGrammar` is false. */
 export interface ThinkingBudgetOptions {
   /** Integer from 0 to 2147483647. 0 forces the end as soon as a block opens, or at once when the prompt ends inside one. */
   maxTokens: number;
@@ -251,6 +251,8 @@ export interface ThinkingBudgetOptions {
   endTag: string;
   /** Text forced before `endTag` when the budget runs out; empty by default. */
   forcedMessage?: string;
+  /** `false` keeps `grammar` applied inside a reasoning block, for a grammar that takes the block and `endTag` itself, such as one written for a prompt that ends inside the block; paused, it would resume at its start and take the reply as reasoning. It must accept `forcedMessage` and `endTag` wherever the budget can run out. Omitted, null or true pauses the grammar. Needs the `thinkingBudgetPauseGrammar` capability. */
+  pauseGrammar?: boolean | null;
 }
 
 /** Completion options the loaded core applies. Every flag is false until a model load initializes the core. */
@@ -261,6 +263,8 @@ export interface CompletionCapabilities {
   minP: boolean;
   /** `CompletionOptions.thinkingBudget` is applied. */
   thinkingBudget: boolean;
+  /** `CompletionOptions.thinkingBudget.pauseGrammar` is applied. A core without it pauses every grammar inside a reasoning block. */
+  thinkingBudgetPauseGrammar: boolean;
   /** The penalties count the prompt, and `CompletionOptions.penaltyPromptTokens` is applied. A core without it never counts the prompt, so `penaltyPromptTokens: false` holds there too. */
   penaltyPromptTokens: boolean;
   /** The loaded core applies CompletionOptions.penaltyPromptTokensWithGrammar. Only literal true permits the opt-in. */

js/src/bridge.ts

@@ -2006,7 +2006,7 @@ export class LlamaWebGpuBridge implements PublicLlamaWebGpuBridge {
     options: BridgeCompletionOptions = {},
   ): Promise<string> {
     const sampling = resolveCompletionSamplingOptions(options);
-    if (sampling.penaltyPromptTokensWithGrammar) {
+    if (sampling.penaltyPromptTokensWithGrammar || sampling.thinkingBudget?.pauseGrammar === false) {
       // Probe the actual worker too: an older worker wrapper may ignore options.
       requireCompletionCapabilities(sampling, await this._getCompletionCapabilitiesUnlocked());
     }