From e5eb15918ec3fec50fb0ebbb45e95a93f141bb10 Mon Sep 17 00:00:00 2001
From: rcourtman <courtmanr@gmail.com>
Date: Tue, 3 Feb 2026 13:12:17 +0000
Subject: [PATCH] Sanitize LLM control tokens from OpenAI-compatible responses

Some local models (llama.cpp, LM Studio) output internal control tokens
like <|channel|>, <|constrain|>, <|message|> instead of using proper
function calling. These tokens leak into the UI creating a poor UX.

This adds sanitization to strip these control tokens from both streaming
and non-streaming responses before they reach the user.
---
 internal/ai/providers/openai.go | 27 ++++++++++++++++++++++-----
 1 file changed, 22 insertions(+), 5 deletions(-)

diff --git a/internal/ai/providers/openai.go b/internal/ai/providers/openai.go
index 80b760da4..7035afe23 100644
--- a/internal/ai/providers/openai.go
+++ b/internal/ai/providers/openai.go
@@ -8,12 +8,23 @@ import (
 	"io"
 	"net/http"
 	"net/url"
+	"regexp"
 	"strings"
 	"time"
 
 	"github.com/rs/zerolog/log"
 )
 
+// controlTokenPattern matches common LLM control tokens that shouldn't appear in output
+// These are internal tokens used by some models (llama.cpp, etc.) for chat templating
+var controlTokenPattern = regexp.MustCompile(`<\|[a-zA-Z_]+\|>`)
+
+// sanitizeContent removes LLM control tokens from content that shouldn't be shown to users
+func sanitizeContent(content string) string {
+	// Remove control tokens like <|channel|>, <|constrain|>, <|message|>, etc.
+	return controlTokenPattern.ReplaceAllString(content, "")
+}
+
 const (
 	openaiAPIURL         = "https://api.openai.com/v1/chat/completions"
 	openaiMaxRetries     = 3
@@ -425,6 +436,9 @@ func (c *OpenAIClient) Chat(ctx context.Context, req ChatRequest) (*ChatResponse
 		contentToUse = choice.Message.ReasoningContent
 	}
 
+	// Sanitize content to remove any leaked control tokens from local models
+	contentToUse = sanitizeContent(contentToUse)
+
 	result := &ChatResponse{
 		Content:          contentToUse,
 		ReasoningContent: choice.Message.ReasoningContent, // DeepSeek thinking mode
@@ -774,12 +788,15 @@ func (c *OpenAIClient) ChatStream(ctx context.Context, req ChatRequest, callback
 
 				delta := choice.Delta
 
-				// Regular content
+				// Regular content - sanitize to remove any leaked control tokens
 				if delta.Content != "" {
-					callback(StreamEvent{
-						Type: "content",
-						Data: ContentEvent{Text: delta.Content},
-					})
+					sanitized := sanitizeContent(delta.Content)
+					if sanitized != "" {
+						callback(StreamEvent{
+							Type: "content",
+							Data: ContentEvent{Text: sanitized},
+						})
+					}
 				}
 
 				// Reasoning content (DeepSeek)