lobe-chat

1
import { afterEach, describe, expect, it, vi } from 'vitest';
2

3
import { MESSAGE_CANCEL_FLAT } from '@/const/message';
4
import { ChatMessageError } from '@/types/message';
5
import { sleep } from '@/utils/sleep';
6

7
import { FetchEventSourceInit } from '../fetchEventSource';
8
import { fetchEventSource } from '../fetchEventSource';
9
import { fetchSSE } from '../fetchSSE';
10

11
// 模拟 i18next
12
vi.mock('i18next', () => ({
13
  t: vi.fn((key) => `translated_${key}`),
14
}));
15

16
vi.mock('../fetchEventSource', () => ({
17
  fetchEventSource: vi.fn(),
18
}));
19

20
// 在每次测试后清理所有模拟
21
afterEach(() => {
22
  vi.restoreAllMocks();
23
});
24

25
describe('fetchSSE', () => {
26
  it('should handle text event correctly', async () => {
27
    const mockOnMessageHandle = vi.fn();
28
    const mockOnFinish = vi.fn();
29

30
    (fetchEventSource as any).mockImplementationOnce(
31
      (url: string, options: FetchEventSourceInit) => {
32
        options.onopen!({ clone: () => ({ ok: true, headers: new Headers() }) } as any);
33
        options.onmessage!({ event: 'text', data: JSON.stringify('Hello') } as any);
34
        options.onmessage!({ event: 'text', data: JSON.stringify(' World') } as any);
35
      },
36
    );
37

38
    await fetchSSE('/', {
39
      onMessageHandle: mockOnMessageHandle,
40
      onFinish: mockOnFinish,
41
      smoothing: false,
42
    });
43

44
    expect(mockOnMessageHandle).toHaveBeenNthCalledWith(1, { text: 'Hello', type: 'text' });
45
    expect(mockOnMessageHandle).toHaveBeenNthCalledWith(2, { text: ' World', type: 'text' });
46
    expect(mockOnFinish).toHaveBeenCalledWith('Hello World', {
47
      observationId: null,
48
      toolCalls: undefined,
49
      traceId: null,
50
      type: 'done',
51
    });
52
  });
53

54
  it('should handle tool_calls event correctly', async () => {
55
    const mockOnMessageHandle = vi.fn();
56
    const mockOnFinish = vi.fn();
57

58
    (fetchEventSource as any).mockImplementationOnce(
59
      (url: string, options: FetchEventSourceInit) => {
60
        options.onopen!({ clone: () => ({ ok: true, headers: new Headers() }) } as any);
61
        options.onmessage!({
62
          event: 'tool_calls',
63
          data: JSON.stringify([
64
            { index: 0, id: '1', type: 'function', function: { name: 'func1', arguments: 'arg1' } },
65
          ]),
66
        } as any);
67
        options.onmessage!({
68
          event: 'tool_calls',
69
          data: JSON.stringify([
70
            { index: 1, id: '2', type: 'function', function: { name: 'func2', arguments: 'arg2' } },
71
          ]),
72
        } as any);
73
      },
74
    );
75

76
    await fetchSSE('/', {
77
      onMessageHandle: mockOnMessageHandle,
78
      onFinish: mockOnFinish,
79
      smoothing: false,
80
    });
81

82
    expect(mockOnMessageHandle).toHaveBeenNthCalledWith(1, {
83
      tool_calls: [{ id: '1', type: 'function', function: { name: 'func1', arguments: 'arg1' } }],
84
      type: 'tool_calls',
85
    });
86
    expect(mockOnMessageHandle).toHaveBeenNthCalledWith(2, {
87
      tool_calls: [
88
        { id: '1', type: 'function', function: { name: 'func1', arguments: 'arg1' } },
89
        { id: '2', type: 'function', function: { name: 'func2', arguments: 'arg2' } },
90
      ],
91
      type: 'tool_calls',
92
    });
93
    expect(mockOnFinish).toHaveBeenCalledWith('', {
94
      observationId: null,
95
      toolCalls: [
96
        { id: '1', type: 'function', function: { name: 'func1', arguments: 'arg1' } },
97
        { id: '2', type: 'function', function: { name: 'func2', arguments: 'arg2' } },
98
      ],
99
      traceId: null,
100
      type: 'done',
101
    });
102
  });
103

104
  it('should call onMessageHandle with full text if no message event', async () => {
105
    const mockOnMessageHandle = vi.fn();
106
    const mockOnFinish = vi.fn();
107

108
    (fetchEventSource as any).mockImplementationOnce(
109
      (url: string, options: FetchEventSourceInit) => {
110
        const res = new Response('Hello World', { status: 200, statusText: 'OK' });
111
        options.onopen!(res as any);
112
      },
113
    );
114

115
    await fetchSSE('/', { onMessageHandle: mockOnMessageHandle, onFinish: mockOnFinish });
116

117
    expect(mockOnMessageHandle).toHaveBeenCalledWith({ text: 'Hello World', type: 'text' });
118
    expect(mockOnFinish).toHaveBeenCalledWith('Hello World', {
119
      observationId: null,
120
      toolCalls: undefined,
121
      traceId: null,
122
      type: 'done',
123
    });
124
  });
125

126
  it('should handle text event with smoothing correctly', async () => {
127
    const mockOnMessageHandle = vi.fn();
128
    const mockOnFinish = vi.fn();
129

130
    (fetchEventSource as any).mockImplementationOnce(
131
      async (url: string, options: FetchEventSourceInit) => {
132
        options.onopen!({ clone: () => ({ ok: true, headers: new Headers() }) } as any);
133
        options.onmessage!({ event: 'text', data: JSON.stringify('Hello') } as any);
134
        await sleep(100);
135
        options.onmessage!({ event: 'text', data: JSON.stringify(' World') } as any);
136
      },
137
    );
138

139
    await fetchSSE('/', {
140
      onMessageHandle: mockOnMessageHandle,
141
      onFinish: mockOnFinish,
142
      smoothing: true,
143
    });
144

145
    expect(mockOnMessageHandle).toHaveBeenNthCalledWith(1, { text: 'Hell', type: 'text' });
146
    expect(mockOnMessageHandle).toHaveBeenNthCalledWith(2, { text: 'o', type: 'text' });
147
    expect(mockOnMessageHandle).toHaveBeenNthCalledWith(3, { text: ' World', type: 'text' });
148
    // more assertions for each character...
149
    expect(mockOnFinish).toHaveBeenCalledWith('Hello World', {
150
      observationId: null,
151
      toolCalls: undefined,
152
      traceId: null,
153
      type: 'done',
154
    });
155
  });
156

157
  it('should handle tool_calls event with smoothing correctly', async () => {
158
    const mockOnMessageHandle = vi.fn();
159
    const mockOnFinish = vi.fn();
160

161
    (fetchEventSource as any).mockImplementationOnce(
162
      (url: string, options: FetchEventSourceInit) => {
163
        options.onopen!({ clone: () => ({ ok: true, headers: new Headers() }) } as any);
164
        options.onmessage!({
165
          event: 'tool_calls',
166
          data: JSON.stringify([
167
            { index: 0, id: '1', type: 'function', function: { name: 'func1', arguments: 'a' } },
168
          ]),
169
        } as any);
170
        options.onmessage!({
171
          event: 'tool_calls',
172
          data: JSON.stringify([
173
            { index: 0, function: { arguments: 'rg1' } },
174
            { index: 1, id: '2', type: 'function', function: { name: 'func2', arguments: 'a' } },
175
          ]),
176
        } as any);
177
        options.onmessage!({
178
          event: 'tool_calls',
179
          data: JSON.stringify([{ index: 1, function: { arguments: 'rg2' } }]),
180
        } as any);
181
      },
182
    );
183

184
    await fetchSSE('/', {
185
      onMessageHandle: mockOnMessageHandle,
186
      onFinish: mockOnFinish,
187
      smoothing: true,
188
    });
189

190
    // TODO: need to check whether the `aarg1` is correct
191
    expect(mockOnMessageHandle).toHaveBeenNthCalledWith(1, {
192
      isAnimationActives: [true, true],
193
      tool_calls: [
194
        { id: '1', type: 'function', function: { name: 'func1', arguments: 'aarg1' } },
195
        { function: { arguments: 'aarg2', name: 'func2' }, id: '2', type: 'function' },
196
      ],
197
      type: 'tool_calls',
198
    });
199
    expect(mockOnMessageHandle).toHaveBeenNthCalledWith(2, {
200
      isAnimationActives: [true, true],
201
      tool_calls: [
202
        { id: '1', type: 'function', function: { name: 'func1', arguments: 'aarg1' } },
203
        { id: '2', type: 'function', function: { name: 'func2', arguments: 'aarg2' } },
204
      ],
205
      type: 'tool_calls',
206
    });
207

208
    // more assertions for each character...
209
    expect(mockOnFinish).toHaveBeenCalledWith('', {
210
      observationId: null,
211
      toolCalls: [
212
        { id: '1', type: 'function', function: { name: 'func1', arguments: 'arg1' } },
213
        { id: '2', type: 'function', function: { name: 'func2', arguments: 'arg2' } },
214
      ],
215
      traceId: null,
216
      type: 'done',
217
    });
218
  });
219

220
  it('should handle request interruption and resumption correctly', async () => {
221
    const mockOnMessageHandle = vi.fn();
222
    const mockOnFinish = vi.fn();
223
    const abortController = new AbortController();
224

225
    (fetchEventSource as any).mockImplementationOnce(
226
      async (url: string, options: FetchEventSourceInit) => {
227
        options.onopen!({ clone: () => ({ ok: true, headers: new Headers() }) } as any);
228
        options.onmessage!({ event: 'text', data: JSON.stringify('Hello') } as any);
229
        await sleep(100);
230
        abortController.abort();
231
        options.onmessage!({ event: 'text', data: JSON.stringify(' World') } as any);
232
      },
233
    );
234

235
    await fetchSSE('/', {
236
      onMessageHandle: mockOnMessageHandle,
237
      onFinish: mockOnFinish,
238
      signal: abortController.signal,
239
      smoothing: true,
240
    });
241

242
    expect(mockOnMessageHandle).toHaveBeenNthCalledWith(1, { text: 'Hell', type: 'text' });
243
    expect(mockOnMessageHandle).toHaveBeenNthCalledWith(2, { text: 'o', type: 'text' });
244
    expect(mockOnMessageHandle).toHaveBeenNthCalledWith(3, { text: ' World', type: 'text' });
245

246
    expect(mockOnFinish).toHaveBeenCalledWith('Hello World', {
247
      type: 'done',
248
      observationId: null,
249
      traceId: null,
250
    });
251
  });
252

253
  it('should call onFinish with correct parameters for different finish types', async () => {
254
    const mockOnFinish = vi.fn();
255

256
    (fetchEventSource as any).mockImplementationOnce(
257
      (url: string, options: FetchEventSourceInit) => {
258
        options.onopen!({ clone: () => ({ ok: true, headers: new Headers() }) } as any);
259
        options.onmessage!({ event: 'text', data: JSON.stringify('Hello') } as any);
260
        options.onerror!({ name: 'AbortError' });
261
      },
262
    );
263

264
    await fetchSSE('/', { onFinish: mockOnFinish, smoothing: false });
265

266
    expect(mockOnFinish).toHaveBeenCalledWith('Hello', {
267
      observationId: null,
268
      toolCalls: undefined,
269
      traceId: null,
270
      type: 'abort',
271
    });
272

273
    (fetchEventSource as any).mockImplementationOnce(
274
      (url: string, options: FetchEventSourceInit) => {
275
        options.onopen!({ clone: () => ({ ok: true, headers: new Headers() }) } as any);
276
        options.onmessage!({ event: 'text', data: JSON.stringify('Hello') } as any);
277
        options.onerror!(new Error('Unknown error'));
278
      },
279
    );
280

281
    await fetchSSE('/', { onFinish: mockOnFinish, smoothing: false });
282

283
    expect(mockOnFinish).toHaveBeenCalledWith('Hello', {
284
      observationId: null,
285
      toolCalls: undefined,
286
      traceId: null,
287
      type: 'error',
288
    });
289
  });
290

291
  describe('onAbort', () => {
292
    it('should call onAbort when AbortError is thrown', async () => {
293
      const mockOnAbort = vi.fn();
294

295
      (fetchEventSource as any).mockImplementationOnce(
296
        (url: string, options: FetchEventSourceInit) => {
297
          options.onmessage!({ event: 'text', data: JSON.stringify('Hello') } as any);
298
          options.onerror!({ name: 'AbortError' });
299
        },
300
      );
301

302
      await fetchSSE('/', { onAbort: mockOnAbort, smoothing: false });
303

304
      expect(mockOnAbort).toHaveBeenCalledWith('Hello');
305
    });
306

307
    it('should call onAbort when MESSAGE_CANCEL_FLAT is thrown', async () => {
308
      const mockOnAbort = vi.fn();
309

310
      (fetchEventSource as any).mockImplementationOnce(
311
        (url: string, options: FetchEventSourceInit) => {
312
          options.onmessage!({ event: 'text', data: JSON.stringify('Hello') } as any);
313
          options.onerror!(MESSAGE_CANCEL_FLAT);
314
        },
315
      );
316

317
      await fetchSSE('/', { onAbort: mockOnAbort, smoothing: false });
318

319
      expect(mockOnAbort).toHaveBeenCalledWith('Hello');
320
    });
321
  });
322

323
  describe('onErrorHandle', () => {
324
    it('should call onErrorHandle when Chat Message error is thrown', async () => {
325
      const mockOnErrorHandle = vi.fn();
326
      const mockError: ChatMessageError = {
327
        body: {},
328
        message: 'StreamChunkError',
329
        type: 'StreamChunkError',
330
      };
331

332
      (fetchEventSource as any).mockImplementationOnce(
333
        (url: string, options: FetchEventSourceInit) => {
334
          options.onerror!(mockError);
335
        },
336
      );
337

338
      try {
339
        await fetchSSE('/', { onErrorHandle: mockOnErrorHandle });
340
      } catch (e) {}
341

342
      expect(mockOnErrorHandle).toHaveBeenCalledWith(mockError);
343
    });
344

345
    it('should call onErrorHandle when Unknown error is thrown', async () => {
346
      const mockOnErrorHandle = vi.fn();
347
      const mockError = new Error('Unknown error');
348

349
      (fetchEventSource as any).mockImplementationOnce(
350
        (url: string, options: FetchEventSourceInit) => {
351
          options.onerror!(mockError);
352
        },
353
      );
354

355
      try {
356
        await fetchSSE('/', { onErrorHandle: mockOnErrorHandle });
357
      } catch (e) {}
358

359
      expect(mockOnErrorHandle).toHaveBeenCalledWith({
360
        type: 'UnknownChatFetchError',
361
        message: 'Unknown error',
362
        body: {
363
          message: 'Unknown error',
364
          name: 'Error',
365
          stack: expect.any(String),
366
        },
367
      });
368
    });
369

370
    it('should call onErrorHandle when response is not ok', async () => {
371
      const mockOnErrorHandle = vi.fn();
372

373
      (fetchEventSource as any).mockImplementationOnce(
374
        async (url: string, options: FetchEventSourceInit) => {
375
          const res = new Response(JSON.stringify({ errorType: 'SomeError' }), {
376
            status: 400,
377
            statusText: 'Error',
378
          });
379

380
          try {
381
            await options.onopen!(res as any);
382
          } catch (e) {}
383
        },
384
      );
385

386
      try {
387
        await fetchSSE('/', { onErrorHandle: mockOnErrorHandle });
388
      } catch (e) {
389
        expect(mockOnErrorHandle).toHaveBeenCalledWith({
390
          body: undefined,
391
          message: 'translated_response.SomeError',
392
          type: 'SomeError',
393
        });
394
      }
395
    });
396

397
    it('should call onErrorHandle when stream chunk has error type', async () => {
398
      const mockOnErrorHandle = vi.fn();
399
      const mockError = {
400
        type: 'StreamChunkError',
401
        message: 'abc',
402
        body: { message: 'abc', context: {} },
403
      };
404

405
      (fetchEventSource as any).mockImplementationOnce(
406
        (url: string, options: FetchEventSourceInit) => {
407
          options.onmessage!({
408
            event: 'error',
409
            data: JSON.stringify(mockError),
410
          } as any);
411
        },
412
      );
413

414
      try {
415
        await fetchSSE('/', { onErrorHandle: mockOnErrorHandle });
416
      } catch (e) {}
417

418
      expect(mockOnErrorHandle).toHaveBeenCalledWith(mockError);
419
    });
420

421
    it('should call onErrorHandle when stream chunk is not valid json', async () => {
422
      const mockOnErrorHandle = vi.fn();
423
      const mockError = 'abc';
424

425
      (fetchEventSource as any).mockImplementationOnce(
426
        (url: string, options: FetchEventSourceInit) => {
427
          options.onmessage!({ event: 'text', data: mockError } as any);
428
        },
429
      );
430

431
      try {
432
        await fetchSSE('/', { onErrorHandle: mockOnErrorHandle });
433
      } catch (e) {}
434

435
      expect(mockOnErrorHandle).toHaveBeenCalledWith({
436
        body: {
437
          context: {
438
            chunk: 'abc',
439
            error: {
440
              message: 'Unexpected token a in JSON at position 0',
441
              name: 'SyntaxError',
442
            },
443
          },
444
          message:
445
            'chat response streaming chunk parse error, please contact your API Provider to fix it.',
446
        },
447
        message: 'parse error',
448
        type: 'StreamChunkError',
449
      });
450
    });
451
  });
452
});
453