Merge PR #8752: fix(tiny): match title stop string against generated tokens only (@alphastorm)

# Conflicts:
#	packages/coding-agent/test/tiny-title-generator.test.ts

packages/coding-agent/test/tiny-title-generator.test.ts

@@ -1,4 +1,5 @@
 import { afterEach, beforeAll, describe, expect, it, vi } from "bun:test";
+import type { StoppingCriteria, TextGenerationPipeline } from "@huggingface/transformers";
 import type { Api, Model } from "@oh-my-pi/pi-ai";
 import * as ai from "@oh-my-pi/pi-ai";
 import { getBundledModel } from "@oh-my-pi/pi-catalog/models";
@@ -31,6 +32,7 @@ import type { TinyTitleWorkerInbound, TinyTitleWorkerOutbound } from "@oh-my-pi/
 import { generateSessionTitle } from "@oh-my-pi/pi-coding-agent/utils/title-generator";
 import type { Subprocess } from "bun";
 import { buildCompletionPrompt } from "../src/tiny/completion-prompt";
+import { createStopOnTextCriteria, type TransformersRuntime } from "../src/tiny/worker";
 
 function getModelOrThrow(id: string): Model<Api> {
 	const model = getBundledModel("anthropic", id);
@@ -398,3 +400,49 @@ describe("tiny-models CLI", () => {
 		expect(isSubcommand("tiny-models")).toBe(true);
 	});
 });
+
+describe("local title stop criteria", () => {
+	/** Minimal stand-ins: the criteria only needs a StoppingCriteria base to extend
+	 *  and a tokenizer that can decode a token window. */
+	const transformers = { StoppingCriteria: class {} } as unknown as TransformersRuntime;
+	const tokenizer = {
+		decode: (ids: number[]) => ids.map(id => (id === 1 ? "</title>" : "x")).join(""),
+	} as unknown as TextGenerationPipeline["tokenizer"];
+	/** `_call(inputIds, scores)`; the criteria ignores scores. */
+	const call = (criteria: StoppingCriteria, inputIds: number[][]): boolean[] =>
+		criteria._call(
+			inputIds,
+			inputIds.map(() => []),
+		);
+
+	it("ignores a stop string that appears only in the prompt", () => {
+		const criteria = createStopOnTextCriteria(transformers, tokenizer, "</title>");
+		// Token 1 decodes to the stop string and sits inside the prompt.
+		const prompt = [1, 0, 0];
+		expect(call(criteria, [[...prompt, 0]])).toEqual([false]);
+		expect(call(criteria, [[...prompt, 0, 0]])).toEqual([false]);
+	});
+
+	it("stops once the stop string is generated", () => {
+		const criteria = createStopOnTextCriteria(transformers, tokenizer, "</title>");
+		const prompt = [1, 0, 0];
+		expect(call(criteria, [[...prompt, 0]])).toEqual([false]);
+		expect(call(criteria, [[...prompt, 0, 1]])).toEqual([true]);
+	});
+
+	it("tracks each batch entry independently", () => {
+		const criteria = createStopOnTextCriteria(transformers, tokenizer, "</title>");
+		expect(
+			call(criteria, [
+				[1, 0],
+				[0, 0],
+			]),
+		).toEqual([false, false]);
+		expect(
+			call(criteria, [
+				[1, 0, 0],
+				[0, 0, 1],
+			]),
+		).toEqual([false, true]);
+	});
+});