Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 52 additions & 1 deletion web-app/src/containers/MessageItem.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
import { memo, useState, useCallback, useEffect } from 'react'
import type { UIMessage, ChatStatus } from 'ai'
import { RenderMarkdown } from './RenderMarkdown'
import { cn } from '@/lib/utils'
import { cn, getModelDisplayName, getProviderTitle, getProviderLogo } from '@/lib/utils'
import { twMerge } from 'tailwind-merge'
import {
ChainOfThought,
Expand Down Expand Up @@ -87,9 +87,40 @@ export const MessageItem = memo(
onDelete,
}: MessageItemProps) => {
const selectedModel = useModelProvider((state) => state.selectedModel)
const selectedProvider = useModelProvider((state) => state.selectedProvider)
const providers = useModelProvider((state) => state.providers)
const metadata = message.metadata as Record<string, unknown> | undefined
const messageError = useMessageErrors((s) => s.errors[message.id])
const createdAt = (metadata?.createdAt as Date) ?? new Date()

// Derive model display name from per-message metadata
const messageModelId = metadata?.modelId as string | undefined
const messageModelProvider = metadata?.modelProvider as string | undefined

const modelDisplayName = useMemo(() => {
if (messageModelId) {
const provider = providers.find(
(p) => p.provider === messageModelProvider
)
const model = provider?.models.find((m) => m.id === messageModelId)
if (model) return getModelDisplayName(model)
return messageModelId
}
// Backwards compat: for the last message without metadata, use current selection
if (isLastMessage && selectedModel) {
return getModelDisplayName(selectedModel)
}
return null
}, [messageModelId, messageModelProvider, isLastMessage, selectedModel, providers])

// Provider for the logo/badge: per-message metadata, falling back to the
// current selection for the streaming last message before metadata persists.
const modelProviderForDisplay =
messageModelProvider ?? (isLastMessage ? selectedProvider : undefined)
const modelProviderLogo = modelProviderForDisplay
? getProviderLogo(modelProviderForDisplay)
: undefined

const [previewImage, setPreviewImage] = useState<{
url: string
filename?: string
Expand Down Expand Up @@ -549,6 +580,26 @@ export const MessageItem = memo(
)}
>

{/* Model name label for assistant messages */}
{message.role === 'assistant' && modelDisplayName && (
<div className="flex items-center gap-1.5 mb-1.5 text-xs text-muted-foreground">
{modelProviderLogo ? (
<img
src={modelProviderLogo}
alt={`${getProviderTitle(modelProviderForDisplay!)} logo`}
className="size-4 object-contain rounded-sm"
/>
) : modelProviderForDisplay ? (
<div className="flex size-4 rounded-sm border items-center justify-center">
<span className="text-[9px] leading-0 capitalize">
{getProviderTitle(modelProviderForDisplay).charAt(0)}
</span>
</div>
) : null}
<span className="font-medium">{modelDisplayName}</span>
</div>
)}

{/* Render message parts */}
{renderedParts}

Expand Down
86 changes: 85 additions & 1 deletion web-app/src/containers/__tests__/MessageItem.test.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,15 @@ import '@testing-library/jest-dom'
// ---- Module mocks ----------------------------------------------------------

const selectedModelRef = vi.hoisted(() => ({ current: { id: 'm1' } as any }))
const selectedProviderRef = vi.hoisted(() => ({ current: 'llamacpp' }))
const providersRef = vi.hoisted(() => ({ current: [] as any[] }))
vi.mock('@/hooks/useModelProvider', () => ({
useModelProvider: (selector: any) =>
selector({ selectedModel: selectedModelRef.current }),
selector({
selectedModel: selectedModelRef.current,
selectedProvider: selectedProviderRef.current,
providers: providersRef.current,
}),
}))

// Stub heavy children: RenderMarkdown
Expand Down Expand Up @@ -115,6 +121,8 @@ describe('MessageItem', () => {
beforeEach(() => {
vi.clearAllMocks()
selectedModelRef.current = { id: 'm1' }
selectedProviderRef.current = 'llamacpp'
providersRef.current = []
})

it('renders assistant text via RenderMarkdown', () => {
Expand All @@ -129,6 +137,82 @@ describe('MessageItem', () => {
expect(screen.getByTestId('render-markdown')).toHaveTextContent('Hello assistant')
})

it('renders the model name label and provider logo from message metadata', () => {
providersRef.current = [
{
provider: 'anthropic',
models: [{ id: 'claude-x', displayName: 'Claude X' }],
},
]
render(
<MessageItem
message={
makeMsg({
metadata: {
createdAt: new Date(),
modelId: 'claude-x',
modelProvider: 'anthropic',
},
}) as any
}
isFirstMessage
isLastMessage={false}
status={'ready' as any}
/>
)
expect(screen.getByText('Claude X')).toBeInTheDocument()
expect(screen.getByAltText('Anthropic logo')).toBeInTheDocument()
})

it('falls back to the raw model id when metadata model is unknown', () => {
render(
<MessageItem
message={
makeMsg({
metadata: {
createdAt: new Date(),
modelId: 'mystery-model',
modelProvider: 'openai',
},
}) as any
}
isFirstMessage
isLastMessage={false}
status={'ready' as any}
/>
)
expect(screen.getByText('mystery-model')).toBeInTheDocument()
expect(screen.getByAltText('OpenAI logo')).toBeInTheDocument()
})

it('falls back to current selection for the last message without metadata', () => {
selectedModelRef.current = { id: 'm1', displayName: 'Local Model' }
selectedProviderRef.current = 'llamacpp'
render(
<MessageItem
message={makeMsg() as any}
isFirstMessage
isLastMessage
status={'ready' as any}
/>
)
expect(screen.getByText('Local Model')).toBeInTheDocument()
expect(screen.getByAltText('Llama.cpp logo')).toBeInTheDocument()
})

it('does not render a model label for a non-last message without metadata', () => {
selectedModelRef.current = { id: 'm1', displayName: 'Local Model' }
render(
<MessageItem
message={makeMsg() as any}
isFirstMessage
isLastMessage={false}
status={'ready' as any}
/>
)
expect(screen.queryByText('Local Model')).not.toBeInTheDocument()
})

it('renders user message in a bubble (no markdown renderer)', () => {
render(
<MessageItem
Expand Down
2 changes: 1 addition & 1 deletion web-app/src/index.css
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@
--color-sidebar-ring: var(--sidebar-ring);
--animate-hue-rotate: hue-rotate 4s infinite;
--font-studio: "StudioFeixenSans";
--font-sans: "Inter", sans-serif;
--font-sans: "Inter", "Apple Color Emoji", "Segoe UI Emoji", "Noto Color Emoji", sans-serif;

/* Font scale based on --font-size-base */
--text-xs: calc(var(--font-size-base) * 0.75); /* ~12px */
Expand Down
202 changes: 202 additions & 0 deletions web-app/src/lib/__tests__/sse-event-filter.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,202 @@
import { describe, it, expect } from 'vitest'
import { createSSEEventFilter, createCustomFetch } from '../model-factory'

/**
* Helper: pipe a string through the SSE event filter and collect the output.
*/
async function filterSSE(input: string): Promise<string> {
const encoder = new TextEncoder()
const stream = new ReadableStream<Uint8Array>({
start(controller) {
controller.enqueue(encoder.encode(input))
controller.close()
},
})

const filtered = stream.pipeThrough(createSSEEventFilter())
const reader = filtered.getReader()
const decoder = new TextDecoder()
let result = ''
while (true) {
const { done, value } = await reader.read()
if (done) break
result += decoder.decode(value, { stream: true })
}
return result
}

describe('createSSEEventFilter', () => {
it('passes through standard data blocks (no event: field)', async () => {
const input =
'data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"Hello"}}]}\n\n' +
'data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":" world"}}]}\n\n'
const output = await filterSSE(input)
expect(output).toBe(input)
})

it('passes through data blocks with event: message', async () => {
const input =
'event: message\ndata: {"id":"chatcmpl-1","choices":[{"index":0,"delta":{"content":"Hi"}}]}\n\n'
const output = await filterSSE(input)
expect(output).toBe(input)
})

it('drops data blocks with custom event types', async () => {
const custom =
'event: hermes.tool.progress\ndata: {"tool":"terminal","emoji":"💻","status":"running"}\n\n'
const standard =
'data: {"id":"chatcmpl-1","choices":[{"index":0,"delta":{"content":"Hello"}}]}\n\n'
const output = await filterSSE(custom + standard)
expect(output).toBe(standard)
})

it('drops multiple custom event types while keeping standard ones', async () => {
const custom1 =
'event: hermes.tool.progress\ndata: {"tool":"terminal"}\n\n'
const standard1 =
'data: {"id":"chatcmpl-1","choices":[{"index":0,"delta":{"content":"A"}}]}\n\n'
const custom2 =
'event: heartbeat\ndata: {"ping":true}\n\n'
const standard2 =
'data: {"id":"chatcmpl-1","choices":[{"index":0,"delta":{"content":"B"}}]}\n\n'

const output = await filterSSE(custom1 + standard1 + custom2 + standard2)
expect(output).toBe(standard1 + standard2)
})

it('keeps data: [DONE] blocks', async () => {
const input =
'data: {"id":"chatcmpl-1","choices":[{"index":0,"delta":{}},"finish_reason":"stop"]}\n\n' +
'data: [DONE]\n\n'
const output = await filterSSE(input)
expect(output).toBe(input)
})

it('handles chunks split across multiple stream pieces', async () => {
const encoder = new TextDecoder()
const part1 =
'event: hermes.tool.progress\ndata: {"tool":"term'
const part2 =
'inal","status":"running"}\n\ndata: {"id":"chatcmpl-1","choices":[{"index":0,"delta":{"content":"OK"}}]}\n\n'

const rs = new ReadableStream<Uint8Array>({
start(controller) {
const enc = new TextEncoder()
controller.enqueue(enc.encode(part1))
controller.enqueue(enc.encode(part2))
controller.close()
},
})

const reader = rs.pipeThrough(createSSEEventFilter()).getReader()
let result = ''
while (true) {
const { done, value } = await reader.read()
if (done) break
result += encoder.decode(value, { stream: true })
}

expect(result).toBe(
'data: {"id":"chatcmpl-1","choices":[{"index":0,"delta":{"content":"OK"}}]}\n\n'
)
})

it('handles CRLF line endings', async () => {
const custom =
'event: hermes.tool.progress\r\ndata: {"tool":"terminal"}\r\n\r\n'
const standard =
'data: {"id":"chatcmpl-1","choices":[{"index":0,"delta":{"content":"Hi"}}]}\r\n\r\n'
const output = await filterSSE(custom + standard)
// After CRLF normalization, the filter should have dropped the custom block
// and kept the standard block (with normalized newlines)
expect(output).toContain('"content":"Hi"')
expect(output).not.toContain('hermes.tool.progress')
})

it('handles empty event type (passes through)', async () => {
// An `event:` with an empty value should still pass through
const input = 'event: \ndata: {"id":"chatcmpl-1","choices":[]}\n\n'
const output = await filterSSE(input)
expect(output).toBe(input)
})

it('handles event type with extra whitespace', async () => {
const custom =
'event: hermes.tool.progress \ndata: {"tool":"terminal"}\n\n'
const standard =
'data: {"id":"chatcmpl-1","choices":[{"index":0,"delta":{"content":"OK"}}]}\n\n'
const output = await filterSSE(custom + standard)
expect(output).toBe(standard)
})
})

describe('createCustomFetch — SSE event filtering integration', () => {
function makeSSEFetch(
sseBody: string,
headers?: Record<string, string>
): typeof globalThis.fetch {
return (async () =>
new Response(sseBody, {
status: 200,
headers: {
'content-type': headers?.['content-type'] ?? 'text/event-stream',
...headers,
},
})) as typeof globalThis.fetch
}

it('filters custom events from SSE responses', async () => {
const custom =
'event: hermes.tool.progress\ndata: {"tool":"terminal","status":"running"}\n\n'
const standard =
'data: {"id":"chatcmpl-1","choices":[{"index":0,"delta":{"content":"Hello"}}]}\n\n'
const baseFetch = makeSSEFetch(custom + standard)
const wrapped = createCustomFetch(baseFetch, {}, false)

const res = await wrapped('http://test/v1/chat/completions', {
method: 'POST',
body: '{}',
})

expect(res.ok).toBe(true)
const body = await res.text()
expect(body).not.toContain('hermes.tool.progress')
expect(body).toContain('"content":"Hello"')
})

it('does NOT filter non-SSE responses (e.g. JSON)', async () => {
const jsonBody = JSON.stringify({
id: 'chatcmpl-1',
choices: [{ message: { content: 'Hi' } }],
})
const baseFetch = makeSSEFetch(jsonBody, {
'content-type': 'application/json',
})
const wrapped = createCustomFetch(baseFetch, {}, false)

const res = await wrapped('http://test/v1/chat/completions', {
method: 'POST',
body: '{}',
})

const body = await res.text()
expect(body).toBe(jsonBody)
})

it('passes standard SSE streams through unchanged', async () => {
const standard =
'data: {"id":"chatcmpl-1","choices":[{"index":0,"delta":{"role":"assistant"}}]}\n\n' +
'data: {"id":"chatcmpl-1","choices":[{"index":0,"delta":{"content":"Hello"}}]}\n\n' +
'data: [DONE]\n\n'
const baseFetch = makeSSEFetch(standard)
const wrapped = createCustomFetch(baseFetch, {}, false)

const res = await wrapped('http://test/v1/chat/completions', {
method: 'POST',
body: '{}',
})

const body = await res.text()
expect(body).toBe(standard)
})
})
Loading