1  import { randomUUID } from 'node:crypto'
  2  
  3  import { beforeEach, describe, expect, it, vi } from 'vitest'
  4  
  5  import type { ChatStreamEvent } from '@/lib/shared/chat'
  6  import { streamChatMessageEvents } from '@/server/chat/service'
  7  
  8  const streamReplyMock = vi.hoisted(() => vi.fn())
  9  
 10  vi.mock('@/server/chat/memory-retrieval', () => ({
 11    refreshMemoryEmbeddings: vi.fn(async () => {}),
 12    selectRelevantMemories: vi.fn(async () => []),
 13  }))
 14  
 15  vi.mock('@/server/providers', () => ({
 16    getChatProvider: vi.fn(() => ({
 17      streamReply: streamReplyMock,
 18      generateReply: vi.fn(async () => ({
 19        provider: 'mock-provider',
 20        mocked: false,
 21        text: 'unused',
 22      })),
 23    })),
 24  }))
 25  
 26  describe('streamChatMessageEvents runId and tool lifecycle passthrough', () => {
 27    beforeEach(() => {
 28      vi.clearAllMocks()
 29    })
 30  
 31    it('emits one runId for the full turn and forwards provider tool lifecycle chunks', async () => {
 32      const sessionId = randomUUID()
 33  
 34      streamReplyMock.mockImplementationOnce(async () => ({
 35        provider: 'openai-compatible',
 36        mocked: false,
 37        stream: (async function* () {
 38          yield {
 39            type: 'tool_start' as const,
 40            tool: 'bash',
 41            toolCallId: 'call_1',
 42            source: 'function' as const,
 43            input: { command: 'pwd' },
 44          }
 45          yield {
 46            type: 'tool_progress' as const,
 47            tool: 'bash',
 48            toolCallId: 'call_1',
 49            source: 'function' as const,
 50            progress: 'Executing tool locally...',
 51          }
 52          yield {
 53            type: 'tool_complete' as const,
 54            tool: 'bash',
 55            toolCallId: 'call_1',
 56            source: 'function' as const,
 57            ok: true,
 58            output: 'exit 0 | stdout: /Users/justinedwards/git/helper',
 59          }
 60          yield {
 61            type: 'thinking' as const,
 62            thinking: 'Summarizing tool output...',
 63          }
 64          yield {
 65            type: 'delta' as const,
 66            delta: 'Done.',
 67          }
 68        })(),
 69      }))
 70  
 71      const events: ChatStreamEvent[] = []
 72      for await (const event of streamChatMessageEvents({
 73        sessionId,
 74        text: 'run the command',
 75        dangerousTools: { bash: true },
 76      })) {
 77        events.push(event)
 78      }
 79  
 80      const sessionEvent = events.find(
 81        (event): event is Extract<ChatStreamEvent, { type: 'session' }> =>
 82          event.type === 'session',
 83      )
 84      expect(sessionEvent).toBeDefined()
 85      expect(sessionEvent?.runId).toBeTruthy()
 86  
 87      const runId = sessionEvent?.runId ?? ''
 88      for (const event of events) {
 89        if (event.type === 'session') continue
 90        expect(event.runId).toBe(runId)
 91      }
 92  
 93      expect(
 94        events.some(
 95          (event) =>
 96            event.type === 'tool_start' &&
 97            event.tool === 'bash' &&
 98            event.toolCallId === 'call_1',
 99        ),
100      ).toBe(true)
101      expect(
102        events.some(
103          (event) =>
104            event.type === 'tool_complete' &&
105            event.ok === true &&
106            event.output.includes('stdout'),
107        ),
108      ).toBe(true)
109      expect(
110        events.some(
111          (event) =>
112            event.type === 'thinking' &&
113            event.thinking.includes('Summarizing tool output'),
114        ),
115      ).toBe(true)
116  
117      const doneEvent = events.at(-1)
118      expect(doneEvent?.type).toBe('done')
119      if (doneEvent?.type === 'done') {
120        expect(doneEvent.assistantMessage.text).toBe('Done.')
121      }
122    })
123  })