diff --git a/test/integration/thinking_budget_grammar_test.dart b/test/integration/thinking_budget_grammar_test.dart
index 984687865..f1d06b568 100644
--- a/test/integration/thinking_budget_grammar_test.dart
+++ b/test/integration/thinking_budget_grammar_test.dart
@@ -120,4 +120,49 @@ thought-text ::= [^d]*
);
}
});
+
+ // A lazy grammar is dormant until its trigger, so a budget that leaves it
+ // applied through the thought would lock the reasoning onto its root.
+ test('a lazy grammar stays paused through a spent thinking budget even when '
+ 'its root takes the thought', () async {
+ for (final root in ['think-body', 'seed-thought']) {
+ final grammar =
+ 'root ::= $root "ok"\n'
+ '$root ::= [^<]* ""\n';
+
+ for (final budget in [1, 3]) {
+ final chunks = await backend
+ .generate(
+ context,
+ prompt,
+ GenerationParams(
+ maxTokens: 48,
+ temp: 0,
+ seed: 1,
+ penalty: 1,
+ grammar: grammar,
+ grammarLazy: true,
+ grammarTriggers: const [
+ GenerationGrammarTrigger.typed(
+ type: GrammarTriggerType.word,
+ value: ' a',
+ ),
+ ],
+ thinkingBudget: ThinkingBudget(
+ maxTokens: budget,
+ startTag: '',
+ endTag: '',
+ ),
+ ),
+ )
+ .toList();
+
+ expect(
+ utf8.decode(chunks.expand((chunk) => chunk).toList()),
+ isNot(endsWith('ok')),
+ reason: '$root, budget $budget',
+ );
+ }
+ }
+ });
}
diff --git a/test/unit/core/template/handlers/seed_oss_handler_test.dart b/test/unit/core/template/handlers/seed_oss_handler_test.dart
index 7de91be76..f6ae095a3 100644
--- a/test/unit/core/template/handlers/seed_oss_handler_test.dart
+++ b/test/unit/core/template/handlers/seed_oss_handler_test.dart
@@ -483,7 +483,9 @@ void main() {
expect(
grammar.accepts(
'Plan ${budget.forcedMessage}${budget.endTag}'
- '\n\n',
+ '\n\n'
+ '\nParis\n\n'
+ '\n',
),
isFalse,
);