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/CHANGELOG.md

@@ -43,6 +43,9 @@
 ### Fixed
 
 - Fixed memory extraction sending its instructions, few-shot examples, and the user's message as a single user turn, which caused small local models to echo the examples instead of extracting facts; instructions now travel as a system turn and the raw text as the user turn
+### Fixed
+
+- Fixed local title generation stopping on a stop string that appeared in the prompt instead of the generated tokens
 
 ## [17.3.7] - 2026-08-17
 

packages/coding-agent/src/tiny/worker.ts

@@ -43,7 +43,7 @@ const TINY_TITLE_SYSTEM_PROMPT = prompt.render(titleSystemPrompt);
 const tinyModelDevicePreference = resolveTinyModelDevicePreference();
 const tinyModelDtypeOverride = resolveTinyModelDtypeOverride();
 
-interface TransformersRuntime extends TransformersRuntimeMetadata {
+export interface TransformersRuntime extends TransformersRuntimeMetadata {
 	env: {
 		cacheDir?: string;
 		allowLocalModels?: boolean;
@@ -80,7 +80,13 @@ function getTinyTitleRuntimeDir(): string {
 	);
 }
 
-function createStopOnTextCriteria(
+/** Stops generation at the first occurrence of `text` in the *generated* tokens.
+ *
+ *  The window must be anchored to the generation boundary, not to the end of the
+ *  whole sequence: a prompt that itself contains the stop string (chat-level
+ *  few-shot examples ending in `</title>`, for instance) would otherwise match on
+ *  prompt tokens and stop before the model emits anything. */
+export function createStopOnTextCriteria(
 	transformers: TransformersRuntime,
 	tokenizer: TextGenerationPipeline["tokenizer"],
 	text: string,
@@ -88,6 +94,8 @@ function createStopOnTextCriteria(
 	class StopOnTextCriteria extends transformers.StoppingCriteria {
 		#tokenizer: TextGenerationPipeline["tokenizer"];
 		#text: string;
+		/** First generated index per batch entry, captured on the first call. */
+		#generatedStarts: number[] = [];
 
 		constructor() {
 			super();
@@ -96,8 +104,10 @@ function createStopOnTextCriteria(
 		}
 
 		override _call(inputIds: number[][]): boolean[] {
-			return inputIds.map(ids => {
-				const tail = ids.slice(-STOP_DECODE_WINDOW_TOKENS);
+			return inputIds.map((ids, index) => {
+				const generatedStart = this.#generatedStarts[index] ?? Math.max(0, ids.length - 1);
+				this.#generatedStarts[index] = generatedStart;
+				const tail = ids.slice(Math.max(generatedStart, ids.length - STOP_DECODE_WINDOW_TOKENS));
 				const decoded = this.#tokenizer.decode(tail, {
 					skip_special_tokens: false,
 					clean_up_tokenization_spaces: false,

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]);
+	});
+});