fix: strip leaked special tokens from local model responses (#2284) (#2285)

Local model servers (LM Studio, Ollama) sometimes leak chat template
control tokens like <|im_end|> into visible responses when stop tokens
are not configured on the server side.

Add stripSpecialTokens utility that removes known tokens (ChatML,
Llama 3, Gemma, Phi, Mistral, Qwen, DeepSeek, Command R) via a single
pre-compiled regex. Applied in ThinkBlockStreamer.handleDeepseekChunk
before content is appended to the response.

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Logan Yang 2026-03-12 17:36:31 -07:00 committed by GitHub
parent 0ac60eb31e
commit 6d14b3ee89
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 201 additions and 1 deletions

View file

@ -9,6 +9,7 @@ import {
createAIMessageWithToolCalls,
} from "./nativeToolCalling";
import { logInfo, logWarn } from "@/logger";
import { stripSpecialTokens } from "@/utils/stripSpecialTokens";
/**
* ThinkBlockStreamer handles streaming content from various LLM providers
@ -116,7 +117,7 @@ export class ThinkBlockStreamer {
private handleDeepseekChunk(chunk: any) {
// Handle standard string content
if (typeof chunk.content === "string") {
this.fullResponse += chunk.content;
this.fullResponse += stripSpecialTokens(chunk.content);
}
// Handle deepseek reasoning/thinking content

View file

@ -0,0 +1,131 @@
import { stripSpecialTokens } from "@/utils/stripSpecialTokens";
describe("stripSpecialTokens", () => {
// --- Individual token stripping ---
it("strips ChatML <|im_end|>", () => {
expect(stripSpecialTokens("hello<|im_end|>")).toBe("hello");
});
it("strips ChatML <|im_start|>", () => {
expect(stripSpecialTokens("<|im_start|>assistant")).toBe("assistant");
});
it("strips Llama 3 <|eot_id|>", () => {
expect(stripSpecialTokens("done<|eot_id|>")).toBe("done");
});
it("strips Llama 3 <|start_header_id|>", () => {
expect(stripSpecialTokens("<|start_header_id|>user")).toBe("user");
});
it("strips Llama 3 <|end_header_id|>", () => {
expect(stripSpecialTokens("assistant<|end_header_id|>")).toBe("assistant");
});
it("strips Gemma <end_of_turn>", () => {
expect(stripSpecialTokens("response<end_of_turn>")).toBe("response");
});
it("strips Gemma <start_of_turn>", () => {
expect(stripSpecialTokens("<start_of_turn>model")).toBe("model");
});
it("strips Phi <|end|>", () => {
expect(stripSpecialTokens("text<|end|>")).toBe("text");
});
it("strips Phi <|assistant|>", () => {
expect(stripSpecialTokens("<|assistant|>answer")).toBe("answer");
});
it("strips Phi <|user|>", () => {
expect(stripSpecialTokens("<|user|>question")).toBe("question");
});
it("strips Phi <|system|>", () => {
expect(stripSpecialTokens("<|system|>prompt")).toBe("prompt");
});
it("strips Mistral </s>", () => {
expect(stripSpecialTokens("end</s>")).toBe("end");
});
it("strips Mistral [INST]", () => {
expect(stripSpecialTokens("[INST]input")).toBe("input");
});
it("strips Mistral [/INST]", () => {
expect(stripSpecialTokens("output[/INST]")).toBe("output");
});
it("strips Qwen <|endoftext|>", () => {
expect(stripSpecialTokens("text<|endoftext|>")).toBe("text");
});
it("strips DeepSeek <|end▁of▁sentence|>", () => {
expect(stripSpecialTokens("sentence<|end\u2581of\u2581sentence|>")).toBe("sentence");
});
it("strips Command R <|END_OF_TURN_TOKEN|>", () => {
expect(stripSpecialTokens("turn<|END_OF_TURN_TOKEN|>")).toBe("turn");
});
it("strips Command R <|START_OF_TURN_TOKEN|>", () => {
expect(stripSpecialTokens("<|START_OF_TURN_TOKEN|>next")).toBe("next");
});
// --- Normal text is unchanged ---
it("leaves normal text unchanged", () => {
const text = "This is a perfectly normal response with no special tokens.";
expect(stripSpecialTokens(text)).toBe(text);
});
it("leaves empty string unchanged", () => {
expect(stripSpecialTokens("")).toBe("");
});
it("leaves text with regular angle brackets unchanged", () => {
const html = "<div>Hello <strong>world</strong></div>";
expect(stripSpecialTokens(html)).toBe(html);
});
it("does NOT strip <s> (can appear in normal text)", () => {
// <s> alone is not in the strip list; </s> is (Mistral EOS token) and will be removed
expect(stripSpecialTokens("<s>beginning")).toBe("<s>beginning");
});
it("strips </s> Mistral EOS even when preceded by HTML-looking <s>", () => {
// </s> is always stripped as it is the Mistral end-of-sequence token
expect(stripSpecialTokens("The <s>strikethrough</s> text here.")).toBe(
"The <s>strikethrough text here."
);
});
// --- Mixed content ---
it("strips token at end of real text without affecting the rest", () => {
expect(stripSpecialTokens("Here is the answer.<|im_end|>")).toBe("Here is the answer.");
});
it("strips token at start of real text without affecting the rest", () => {
expect(stripSpecialTokens("<|im_start|>Here is the answer.")).toBe("Here is the answer.");
});
it("strips multiple tokens from a single chunk", () => {
expect(stripSpecialTokens("<|im_start|>assistant\nHello there!<|im_end|>")).toBe(
"assistant\nHello there!"
);
});
it("strips tokens from different families in one pass", () => {
expect(stripSpecialTokens("[INST]question[/INST]<|im_end|>answer<|eot_id|>")).toBe(
"questionanswer"
);
});
it("handles repeated occurrences of the same token", () => {
expect(stripSpecialTokens("<|im_end|>text<|im_end|>more<|im_end|>")).toBe("textmore");
});
});

View file

@ -0,0 +1,68 @@
/**
* Strips known chat template control tokens that local model servers (LM Studio, Ollama)
* sometimes leak into visible responses. These tokens are part of the model's internal
* chat template and should never appear in user-visible output.
*
* Token families covered:
* - ChatML: <|im_end|>, <|im_start|>
* - Llama 3: <|eot_id|>, <|start_header_id|>, <|end_header_id|>
* - Gemma: <end_of_turn>, <start_of_turn>
* - Phi: <|end|>, <|assistant|>, <|user|>, <|system|>
* - Mistral/Generic: </s>, [INST], [/INST]
* - Qwen: <|endoftext|>
* - DeepSeek: <|endofsentence|>
* - Command R: <|END_OF_TURN_TOKEN|>, <|START_OF_TURN_TOKEN|>
*/
/** Known special tokens that must never appear in user-visible output. */
const SPECIAL_TOKEN_PATTERNS: string[] = [
// ChatML
"<|im_end|>",
"<|im_start|>",
// Llama 3
"<|eot_id|>",
"<|start_header_id|>",
"<|end_header_id|>",
// Gemma
"<end_of_turn>",
"<start_of_turn>",
// Phi
"<|end|>",
"<|assistant|>",
"<|user|>",
"<|system|>",
// Mistral/Generic (NOT <s> — can appear in normal text)
"</s>",
"[INST]",
"[/INST]",
// Qwen
"<|endoftext|>",
// DeepSeek (uses Unicode U+2581 LOWER ONE EIGHTH BLOCK for underscores)
"<|end\u2581of\u2581sentence|>",
// Command R
"<|END_OF_TURN_TOKEN|>",
"<|START_OF_TURN_TOKEN|>",
];
/**
* Pre-compiled regex built from all known special token patterns.
* Each token is escaped so that regex metacharacters are treated literally.
* The `g` flag ensures all occurrences in a chunk are removed in one pass.
*/
const SPECIAL_TOKENS_REGEX: RegExp = new RegExp(
SPECIAL_TOKEN_PATTERNS.map((t) => t.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")).join("|"),
"g"
);
/**
* Strips all known chat template control tokens from the provided text.
*
* This is a zero-cost operation when the input contains no special tokens
* (the regex engine scans without allocating a new string).
*
* @param text - Raw text chunk from a local LLM response stream.
* @returns The text with all known special tokens removed.
*/
export function stripSpecialTokens(text: string): string {
return text.replace(SPECIAL_TOKENS_REGEX, "");
}