From b8dbf68613a7291df86a567e5103387821f711a9 Mon Sep 17 00:00:00 2001 From: Sunil Srivatsa Date: Mon, 17 Aug 2026 15:19:06 -0700 Subject: [PATCH] fix(tiny): match title stop string against generated tokens only StopOnTextCriteria decoded the last STOP_DECODE_WINDOW_TOKENS of the whole sequence, so prompt tokens were eligible for matching. A prompt that itself contains the stop string stops generation at the first generated token and yields an empty title. Anchor the window to the generation boundary by recording the first generated index per batch entry. Existing local title models are unaffected: with the assistant-prefill prompt shape, the example `` tags sit outside the 32-token window for normal messages, so no shipping model changes behavior. The bug becomes reachable with any chat-level few-shot prompt that places the stop string near the generation boundary. --- packages/coding-agent/CHANGELOG.md | 4 ++ packages/coding-agent/src/tiny/worker.ts | 18 +++++-- .../test/tiny-title-generator.test.ts | 48 +++++++++++++++++++ 3 files changed, 66 insertions(+), 4 deletions(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 2f7b500ff..053ceb6ef 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### 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 ### Changed diff --git a/packages/coding-agent/src/tiny/worker.ts b/packages/coding-agent/src/tiny/worker.ts index ca056baa6..916b71c6d 100644 --- a/packages/coding-agent/src/tiny/worker.ts +++ b/packages/coding-agent/src/tiny/worker.ts @@ -42,7 +42,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; @@ -79,7 +79,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 ``, 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, @@ -87,6 +93,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(); @@ -95,8 +103,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, diff --git a/packages/coding-agent/test/tiny-title-generator.test.ts b/packages/coding-agent/test/tiny-title-generator.test.ts index 2583e8cbb..a302dcced 100644 --- a/packages/coding-agent/test/tiny-title-generator.test.ts +++ b/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"; @@ -30,6 +31,7 @@ import { import type { TinyTitleWorkerInbound, TinyTitleWorkerOutbound } from "@oh-my-pi/pi-coding-agent/tiny/title-protocol"; import { generateSessionTitle } from "@oh-my-pi/pi-coding-agent/utils/title-generator"; import type { Subprocess } from "bun"; +import { createStopOnTextCriteria, type TransformersRuntime } from "../src/tiny/worker"; function getModelOrThrow(id: string): Model { const model = getBundledModel("anthropic", id); @@ -352,3 +354,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 ? "" : "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, ""); + // 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, ""); + 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, ""); + expect( + call(criteria, [ + [1, 0], + [0, 0], + ]), + ).toEqual([false, false]); + expect( + call(criteria, [ + [1, 0, 0], + [0, 0, 1], + ]), + ).toEqual([false, true]); + }); +});