diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 6e88b6b1f..a47a1514c 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/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 diff --git a/packages/coding-agent/src/tiny/worker.ts b/packages/coding-agent/src/tiny/worker.ts index 56a0c0c4b..13c165e00 100644 --- a/packages/coding-agent/src/tiny/worker.ts +++ b/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 ``, 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, diff --git a/packages/coding-agent/test/tiny-title-generator.test.ts b/packages/coding-agent/test/tiny-title-generator.test.ts index 3e114e582..5a31fbc71 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"; @@ -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 { 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 ? "" : "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]); + }); +});