hermes-agent/apps/desktop/src/components/assistant-ui/streaming.test.tsx

import { AssistantRuntimeProvider, type ThreadMessage, useExternalStoreRuntime } from '@assistant-ui/react'
import { act, fireEvent, render, screen, waitFor } from '@testing-library/react'
import { useEffect, useState } from 'react'
import { beforeEach, describe, expect, it, vi } from 'vitest'

import { Thread } from './thread'

const createdAt = new Date('2026-05-01T00:00:00.000Z')

const resizeObservers = new Set<TestResizeObserver>()

class TestResizeObserver {
  private target: Element | null = null

  constructor(private readonly callback: ResizeObserverCallback) {
    resizeObservers.add(this)
  }

  observe(target: Element) {
    this.target = target
  }

  unobserve() {}

  disconnect() {
    resizeObservers.delete(this)
  }

  trigger(height: number) {
    if (!this.target) {
      return
    }

    this.callback(
      [
        {
          contentRect: { height } as DOMRectReadOnly,
          target: this.target
        } as ResizeObserverEntry
      ],
      this as unknown as ResizeObserver
    )
  }
}

vi.stubGlobal('ResizeObserver', TestResizeObserver)
vi.stubGlobal('requestAnimationFrame', (callback: FrameRequestCallback) =>
  window.setTimeout(() => callback(performance.now()), 0)
)
vi.stubGlobal('cancelAnimationFrame', (id: number) => window.clearTimeout(id))

Element.prototype.scrollTo = function scrollTo() {}

async function wait(ms: number) {
  await act(async () => {
    await new Promise(resolve => window.setTimeout(resolve, ms))
  })
}

function userMessage(): ThreadMessage {
  return {
    id: 'user-1',
    role: 'user',
    content: [{ type: 'text', text: 'Stream a response' }],
    attachments: [],
    createdAt,
    metadata: { custom: {} }
  } as ThreadMessage
}

function assistantMessage(text: string, running = true): ThreadMessage {
  return {
    id: 'assistant-1',
    role: 'assistant',
    content: [{ type: 'text', text }],
    status: running ? { type: 'running' } : { type: 'complete', reason: 'stop' },
    createdAt,
    metadata: {
      unstable_state: null,
      unstable_annotations: [],
      unstable_data: [],
      steps: [],
      custom: {}
    }
  } as ThreadMessage
}

function assistantReasoningMessage(text: string): ThreadMessage {
  return {
    id: 'assistant-reasoning-1',
    role: 'assistant',
    content: [{ type: 'reasoning', text }],
    status: { type: 'complete', reason: 'stop' },
    createdAt,
    metadata: {
      unstable_state: null,
      unstable_annotations: [],
      unstable_data: [],
      steps: [],
      custom: {}
    }
  } as ThreadMessage
}

function assistantMultiReasoningMessage(texts: string[]): ThreadMessage {
  return {
    id: 'assistant-reasoning-multi-1',
    role: 'assistant',
    content: texts.map(text => ({ type: 'reasoning', text })),
    status: { type: 'complete', reason: 'stop' },
    createdAt,
    metadata: {
      unstable_state: null,
      unstable_annotations: [],
      unstable_data: [],
      steps: [],
      custom: {}
    }
  } as ThreadMessage
}

function StreamingHarness() {
  const [messages, setMessages] = useState<ThreadMessage[]>([userMessage()])
  const [isRunning, setIsRunning] = useState(true)

  useEffect(() => {
    const first = window.setTimeout(() => {
      setMessages([userMessage(), assistantMessage('first chunk')])
    }, 50)

    const second = window.setTimeout(() => {
      setMessages([userMessage(), assistantMessage('first chunk second chunk')])
    }, 500)

    const complete = window.setTimeout(() => {
      setMessages([userMessage(), assistantMessage('first chunk second chunk', false)])
      setIsRunning(false)
    }, 700)

    return () => {
      window.clearTimeout(first)
      window.clearTimeout(second)
      window.clearTimeout(complete)
    }
  }, [])

  const runtime = useExternalStoreRuntime<ThreadMessage>({
    messages,
    isRunning,
    onNew: async () => {}
  })

  return (
    <AssistantRuntimeProvider runtime={runtime}>
      <Thread loading={isRunning && messages.at(-1)?.role !== 'assistant' ? 'response' : undefined} />
    </AssistantRuntimeProvider>
  )
}

function ReasoningHarness() {
  const runtime = useExternalStoreRuntime<ThreadMessage>({
    messages: [assistantReasoningMessage(' The user is asking what this file is.')],
    isRunning: false,
    onNew: async () => {}
  })

  return (
    <AssistantRuntimeProvider runtime={runtime}>
      <Thread />
    </AssistantRuntimeProvider>
  )
}

function GroupedReasoningHarness() {
  const runtime = useExternalStoreRuntime<ThreadMessage>({
    messages: [assistantMultiReasoningMessage([' First thought.', ' Second thought.'])],
    isRunning: false,
    onNew: async () => {}
  })

  return (
    <AssistantRuntimeProvider runtime={runtime}>
      <Thread />
    </AssistantRuntimeProvider>
  )
}

describe('assistant-ui streaming renderer', () => {
  beforeEach(() => {
    resizeObservers.clear()
  })

  it('renders assistant text incrementally before completion', async () => {
    const { container } = render(<StreamingHarness />)

    expect(screen.getByRole('status', { name: 'Hermes is loading a response' })).toBeTruthy()

    await wait(80)

    await waitFor(() => {
      expect(container.textContent).toContain('first chunk')
    })
    expect(container.textContent).not.toContain('second chunk')
    expect(screen.queryByRole('status', { name: 'Hermes is loading a response' })).toBeNull()

    await wait(500)

    await waitFor(() => {
      expect(container.textContent).toContain('first chunk second chunk')
    })

    await wait(250)

    await waitFor(() => {
      expect(container.textContent).toContain('first chunk second chunk')
    })
  })

  it('does not pull the viewport back down after the user scrolls up during streaming', async () => {
    const { container } = render(<StreamingHarness />)

    const content = container.querySelector('[data-slot="aui_thread-content"]') as HTMLDivElement
    const viewport = content.parentElement as HTMLDivElement
    let scrollHeight = 1_000

    Object.defineProperty(viewport, 'clientHeight', { configurable: true, value: 200 })
    Object.defineProperty(viewport, 'scrollHeight', {
      configurable: true,
      get: () => scrollHeight
    })

    await wait(80)

    await act(async () => {
      viewport.scrollTop = 800
      fireEvent.scroll(viewport)
    })
    await wait(0)

    await act(async () => {
      fireEvent.wheel(viewport, { deltaY: -120 })
      viewport.scrollTop = 420
      fireEvent.scroll(viewport)
    })

    scrollHeight = 1_200

    await act(async () => {
      for (const observer of resizeObservers) {
        observer.trigger(1_200)
      }
    })
    await wait(0)

    expect(viewport.scrollTop).toBe(420)
  })

  it('renders reasoning text without a leading token space', () => {
    const { container } = render(<ReasoningHarness />)

    fireEvent.click(screen.getByRole('button', { name: /thinking/i }))

    expect(container.querySelector('[data-slot="aui_reasoning-text"]')?.textContent).toBe(
      'The user is asking what this file is.'
    )
  })

  it('groups consecutive reasoning parts under one thinking disclosure', () => {
    const { container } = render(<GroupedReasoningHarness />)

    const disclosures = container.querySelectorAll('[data-slot="tool-block"] > button')
    expect(disclosures.length).toBe(1)

    fireEvent.click(disclosures[0])

    const reasoningParts = container.querySelectorAll('[data-slot="aui_reasoning-text"]')
    expect(reasoningParts.length).toBe(2)
    expect(reasoningParts[0]?.textContent).toBe('First thought.')
    expect(reasoningParts[1]?.textContent).toBe('Second thought.')
  })
})