diff --git a/crates/mesh-llm-ui/src/features/chat/components/MessageRow.test.tsx b/crates/mesh-llm-ui/src/features/chat/components/MessageRow.test.tsx
index 11b6244b03..43b17885f9 100644
--- a/crates/mesh-llm-ui/src/features/chat/components/MessageRow.test.tsx
+++ b/crates/mesh-llm-ui/src/features/chat/components/MessageRow.test.tsx
@@ -253,6 +253,24 @@ describe('MessageRow', () => {
expect(screen.getByText('The capital of France is Paris.')).toBeInTheDocument()
})
+ it('moves Gemma channel thinking into the trace before rendering the final response', () => {
+ render(
+
+ )
+
+ expect(screen.getByText('Thinking trace').closest('[data-thinking-state="complete"]')).toHaveTextContent(
+ 'Check whether the prompt is a test.'
+ )
+ expect(screen.getByText('Test received!').closest('.select-text')).toBeInTheDocument()
+ expect(screen.queryByText(/<\|channel\|>thought/)).not.toBeInTheDocument()
+ expect(screen.queryByText(//)).not.toBeInTheDocument()
+ })
+
it('formats final assistant response text as markdown', () => {
render(
{
])
})
+ it('splits Gemma channel thinking from final response text', () => {
+ expect(
+ splitAssistantThinking('<|channel|>thoughtCheck whether the prompt is a test.Test received!')
+ ).toEqual([
+ { kind: 'thinking', text: 'Check whether the prompt is a test.', open: false },
+ { kind: 'response', text: 'Test received!' }
+ ])
+ })
+
+ it('splits Gemma channel thinking when the thought marker is missing the leading pipe', () => {
+ expect(splitAssistantThinking('thoughtCheck facts.Final answer.')).toEqual([
+ { kind: 'thinking', text: 'Check facts.', open: false },
+ { kind: 'response', text: 'Final answer.' }
+ ])
+ })
+
it('keeps an unclosed think segment open for live streams', () => {
expect(splitAssistantThinking('Checking facts')).toEqual([
{ kind: 'thinking', text: 'Checking facts', open: true }
diff --git a/crates/mesh-llm-ui/src/features/chat/components/thinking-segments.ts b/crates/mesh-llm-ui/src/features/chat/components/thinking-segments.ts
index dfee7697d2..032351d784 100644
--- a/crates/mesh-llm-ui/src/features/chat/components/thinking-segments.ts
+++ b/crates/mesh-llm-ui/src/features/chat/components/thinking-segments.ts
@@ -13,8 +13,15 @@ type SplitAssistantThinkingOptions = {
streaming?: boolean
}
+type TagMatch = {
+ index: number
+ tag: string
+}
+
const THINK_OPEN_TAG = ''
const THINK_CLOSE_TAG = ''
+const GEMMA_THOUGHT_CHANNEL_TAGS = ['<|channel|>thought', 'thought']
+const GEMMA_CHANNEL_BOUNDARY_TAGS = ['<|channel|>', '']
function indexOfTag(value: string, tag: string, fromIndex: number) {
for (let index = fromIndex; index <= value.length - tag.length; index += 1) {
@@ -24,12 +31,67 @@ function indexOfTag(value: string, tag: string, fromIndex: number) {
return -1
}
+function findFirstTag(value: string, tags: string[], fromIndex: number): TagMatch | null {
+ let bestMatch: TagMatch | null = null
+
+ for (const tag of tags) {
+ const index = indexOfTag(value, tag, fromIndex)
+ if (index === -1) continue
+ if (bestMatch === null || index < bestMatch.index) {
+ bestMatch = { index, tag }
+ }
+ }
+
+ return bestMatch
+}
+
+function splitGemmaChannelThinking(body: string): AssistantContentSegment[] | null {
+ if (findFirstTag(body, GEMMA_THOUGHT_CHANNEL_TAGS, 0) === null) return null
+
+ const segments: AssistantContentSegment[] = []
+ let cursor = 0
+
+ while (cursor < body.length) {
+ const open = findFirstTag(body, GEMMA_THOUGHT_CHANNEL_TAGS, cursor)
+ if (open === null) {
+ const responseText = body.slice(cursor)
+ if (responseText.length > 0) {
+ segments.push({ kind: 'response', text: responseText })
+ }
+ break
+ }
+
+ const responseText = body.slice(cursor, open.index)
+ if (responseText.length > 0) {
+ segments.push({ kind: 'response', text: responseText })
+ }
+
+ const thinkingStart = open.index + open.tag.length
+ const nextChannel = findFirstTag(body, GEMMA_CHANNEL_BOUNDARY_TAGS, thinkingStart)
+ if (nextChannel === null) {
+ segments.push({ kind: 'thinking', text: body.slice(thinkingStart), open: true })
+ break
+ }
+
+ const thinkingText = body.slice(thinkingStart, nextChannel.index)
+ if (thinkingText.length > 0) {
+ segments.push({ kind: 'thinking', text: thinkingText, open: false })
+ }
+ cursor = nextChannel.index + nextChannel.tag.length
+ }
+
+ return segments
+}
+
export function splitAssistantThinking(
body: string,
{ streaming = false }: SplitAssistantThinkingOptions = {}
): AssistantContentSegment[] {
if (body.length === 0) return []
+ const gemmaSegments = splitGemmaChannelThinking(body)
+ if (gemmaSegments !== null) return gemmaSegments
+
const segments: AssistantContentSegment[] = []
let cursor = 0
let firstSegment = true