fix(ui): count an answer that ends on its last allowed token as finished

The worker inferred "length" whenever the answer used every allowed token.
Transformers.js also hands the end-of-sequence token to the streamer, so an
answer that ended exactly on the limit was reported as cut and the chat
showed "Shortened (token limit)". The worker now remembers the last token
and reports "length" only when it is not an end-of-sequence id of the
model or tokenizer.

Signed-off-by: Bernhard Jackiewicz <bj@techport.io>

graph-ui/src/browser-ai/browser-ai.worker.test.ts

@@ -22,8 +22,8 @@ vi.mock('@huggingface/transformers', () => ({
     },
     TextStreamer: class {
         callback: (chunk: string) => void;
-        token: () => void;
-        constructor(_tokenizer: unknown, options: { callback_function: (chunk: string) => void; token_callback_function: () => void }) {
+        token: (tokens?: bigint[]) => void;
+        constructor(_tokenizer: unknown, options: { callback_function: (chunk: string) => void; token_callback_function: (tokens?: bigint[]) => void }) {
             this.callback = options.callback_function; this.token = options.token_callback_function;
         }
     },
@@ -133,10 +133,11 @@ describe('browser generation worker', () => {
     });
 
     it('reports whether an answer ended at the output limit, at the end of sequence, or on Stop', async () => {
+        fake.modelLoad.mockResolvedValue({ generate: fake.generate, generation_config: { eos_token_id: [2, 3] } });
         await send({ id: 1, kind: 'prepare' });
         const messages = [{ role: 'user' as const, content: 'List' }];
-        const emit = (tokens: number) => async ({ streamer }: { streamer: { callback: (chunk: string) => void; token: () => void } }) => {
-            for (let index = 0; index < tokens; index++) { streamer.token(); streamer.callback('x'); }
+        const emit = (tokens: number, last = 7n) => async ({ streamer }: { streamer: { callback: (chunk: string) => void; token: (tokens?: bigint[]) => void } }) => {
+            for (let index = 0; index < tokens; index++) { streamer.token([index + 1 === tokens ? last : 9n]); streamer.callback('x'); }
         };
         fake.generate.mockImplementationOnce(emit(64));
         await send({ id: 2, kind: 'chat', messages, maxOutputTokens: 64 });
@@ -144,6 +145,10 @@ describe('browser generation worker', () => {
         fake.generate.mockImplementationOnce(emit(10));
         await send({ id: 3, kind: 'chat', messages, maxOutputTokens: 64 });
         expect(replies().at(-1)).toEqual({ id: 3, kind: 'answer', output: 'x'.repeat(10), stopReason: 'eos' });
+        // The streamer also receives the end-of-sequence token; ending on the last allowed token is not a cut.
+        fake.generate.mockImplementationOnce(emit(64, 3n));
+        await send({ id: 4, kind: 'chat', messages, maxOutputTokens: 64 });
+        expect(replies().at(-1)).toEqual({ id: 4, kind: 'answer', output: 'x'.repeat(64), stopReason: 'eos' });
     });
 
     it('rejects unknown profiles and oversized automatic answers without invalidating the model', async () => {

graph-ui/src/browser-ai/browser-ai.worker.ts

@@ -158,11 +158,11 @@ async function generate(id: number, messages: readonly BrowserChatMessage[], opt
         throw new Error(`This conversation uses ${count.toLocaleString()} input tokens. The browser limit is ${selected.contextTokens.toLocaleString()}, including ${outputTokens} reserved for the answer. Remove earlier messages or attach a smaller selection; no code was truncated.`);
     }
     stopping.reset();
-    let answer = '', generated = 0;
+    let answer = '', generated = 0, last: bigint | undefined;
     const streamer = new TextStreamer(tokenizer!, {
         skip_prompt: true,
         callback_function: chunk => { answer += chunk; post({ id, kind: 'token', output: chunk }); },
-        token_callback_function: () => { generated += 1; },
+        token_callback_function: tokens => { generated += 1; last = tokens?.at(-1); },
     });
     const controls = REPETITION_CONTROLS[automaticExplanation ? 'automatic-explanation' : 'chat'];
     const logitsProcessor = new LogitsProcessorList();
@@ -181,8 +181,11 @@ async function generate(id: number, messages: readonly BrowserChatMessage[], opt
     }
     if (fatalFailure) throw fatalFailure;
     if (!answer.trim() && !stopping.interrupted) throw new Error('The model returned no answer.');
-    // A stopped answer is reported as stopped; a full output budget means the answer was cut.
-    return { output: answer, stopReason: stopping.interrupted ? 'interrupted' : generated >= outputTokens ? 'length' : 'eos' };
+    // A stopped answer is reported as stopped; a full output budget means the answer was cut,
+    // unless its last token was the end of sequence (the streamer also receives that token).
+    const ends = [model!.generation_config?.eos_token_id ?? [], tokenizer!.eos_token_id ?? []].flat().map(token => BigInt(token));
+    const ended = last !== undefined && ends.includes(last);
+    return { output: answer, stopReason: stopping.interrupted ? 'interrupted' : generated >= outputTokens && !ended ? 'length' : 'eos' };
 }
 
 self.onmessage = async (event: MessageEvent<BrowserWorkerRequest>) => {