fix: gemma thinking

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(
+      <MessageRow
+        messageRole="assistant"
+        body="<|channel|>thoughtCheck whether the prompt is a test.<channel|>Test received!"
+        timestamp="12:11"
+        model="unsloth/gemma-4-31B-it-GGUF:UD-Q8_K_XL"
+      />
+    )
+
+    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(/<channel\|>/)).not.toBeInTheDocument()
+  })
+
   it('formats final assistant response text as markdown', () => {
     render(
       <MessageRow

crates/mesh-llm-ui/src/features/chat/components/thinking-segments.test.ts

@@ -22,6 +22,22 @@ describe('splitAssistantThinking', () => {
     ])
   })
 
+  it('splits Gemma channel thinking from final response text', () => {
+    expect(
+      splitAssistantThinking('<|channel|>thoughtCheck whether the prompt is a test.<channel|>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('<channel|>thoughtCheck facts.<channel|>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('<think>Checking facts')).toEqual([
       { kind: 'thinking', text: 'Checking facts', open: true }

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 = '<think>'
 const THINK_CLOSE_TAG = '</think>'
+const GEMMA_THOUGHT_CHANNEL_TAGS = ['<|channel|>thought', '<channel|>thought']
+const GEMMA_CHANNEL_BOUNDARY_TAGS = ['<|channel|>', '<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