RAG/client/tests/api.test.ts

import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import { chatRequest, chatSSERequest, clearApiConfig, setApiConfig } from '../src/api';
import { parseConfig } from '../src/config';
import type { RagSource } from '../src/types';

const source: RagSource = {
  documentId: '1234567890123456789', title: 'manual', sourceName: null,
  chunkIndex: 0, score: null, snippet: '实际命中的文档',
};
const encoder = new TextEncoder();
const chunk = (content: string) => `data: ${JSON.stringify({ choices: [{ delta: { content } }] })}\r\n\r\n`;
const metadata = `data: ${JSON.stringify({ object: 'chat.completion.chunk', choices: [], sources: [source] })}\r\n\r\n`;

function response(parts: Uint8Array[], close = true) {
  const cancel = vi.fn();
  const stream = new ReadableStream<Uint8Array>({
    start(controller) {
      parts.forEach(part => controller.enqueue(part));
      if (close) controller.close();
    },
    cancel,
  });
  vi.stubGlobal('fetch', vi.fn().mockResolvedValue(new Response(stream)));
  return { stream, cancel };
}
function receive(signal?: AbortSignal) {
  const onChunk = vi.fn();
  const onDone = vi.fn();
  const onError = vi.fn();
  const onSources = vi.fn();
  const pending = chatSSERequest('问题', onChunk, onDone, onError, 7, true, undefined, signal, onSources);
  return { pending, onChunk, onDone, onError, onSources };
}

beforeEach(() => {
  const config = parseConfig({ integrateId: '42', requestDomain: 'https://example.test', userId: 'user' })!;
  config.chatId = 'conversation';
  setApiConfig(config);
});
afterEach(() => { clearApiConfig(); vi.unstubAllGlobals(); });

describe('same-request answer transport', () => {
  it('stops on DONE without EOF, preserving split UTF-8, CRLF and sources outside text', async () => {
    const bytes = encoder.encode(chunk('中文 ') + metadata + 'data: [DONE]\r\n\r\n' + chunk('不得出现'));
    const { stream, cancel } = response(Array.from(bytes, byte => Uint8Array.of(byte)), false);
    const result = receive();
    await result.pending;
    expect(result.onChunk.mock.calls.flat().join('')).toBe('中文 ');
    expect(result.onSources).toHaveBeenCalledExactlyOnceWith([source]);
    expect(result.onDone).toHaveBeenCalledTimes(1);
    expect(result.onError).not.toHaveBeenCalled();
    expect(cancel).toHaveBeenCalledTimes(1);
    expect(stream.locked).toBe(false);
    expect(fetch).toHaveBeenCalledTimes(1);
    const url = new URL(vi.mocked(fetch).mock.calls[0]![0] as string);
    expect(url.pathname).toBe('/ai/chat/stream');
    expect(url.searchParams.get('rewriteStrategy')).toBe('NONE');
  });

  it('keeps multiline data whitespace and flushes the last unterminated event at EOF', async () => {
    response([encoder.encode('event: status\ndata: 不显示\n\ndata:   first  \r\ndata: second\r\n\r\ndata: last  ')]);
    const result = receive();
    await result.pending;
    expect(result.onChunk.mock.calls.flat()).toEqual(['  first  \nsecond', 'last  ']);
    expect(result.onDone).toHaveBeenCalledTimes(1);
  });

  it('recognizes EOF remainder DONE rather than rendering it', async () => {
    response([encoder.encode(chunk('answer') + 'data: [DONE]')]);
    const result = receive();
    await result.pending;
    expect(result.onChunk).toHaveBeenCalledExactlyOnceWith('answer');
    expect(result.onDone).toHaveBeenCalledTimes(1);
  });

  it('delivers original server error once even when reader cleanup rejects', async () => {
    const { cancel, stream } = response([encoder.encode('data: {"error":{"message":"权限不足"}}\n\n')], false);
    cancel.mockRejectedValue(new Error('cleanup failed'));
    const result = receive();
    await result.pending;
    expect(result.onError).toHaveBeenCalledTimes(1);
    expect(result.onError.mock.calls[0]![0].message).toBe('权限不足');
    expect(result.onDone).not.toHaveBeenCalled();
    expect(stream.locked).toBe(false);
  });

  it('cancels a pending read once without consuming later metadata', async () => {
    const { stream } = response([], false);
    const controller = new AbortController();
    const result = receive(controller.signal);
    await Promise.resolve();
    await Promise.resolve();
    controller.abort();
    await result.pending;
    expect(result.onDone).toHaveBeenCalledTimes(1);
    expect(result.onError).not.toHaveBeenCalled();
    expect(result.onSources).not.toHaveBeenCalled();
    expect(stream.locked).toBe(false);
  });

  it('uses one synchronous JSON request with the same RAG, role and category parameters', async () => {
    const answer = { text: 'answer', mcpEvents: [], suggestions: [], sources: [source] };
    vi.stubGlobal('fetch', vi.fn().mockResolvedValue(Response.json(answer)));
    expect(await chatRequest('问题', ['https://example.test/photo.png'], 7, true)).toEqual(answer);
    expect(fetch).toHaveBeenCalledTimes(1);
    const url = new URL(vi.mocked(fetch).mock.calls[0]![0] as string);
    expect(url.pathname).toBe('/ai/chat/result');
    expect(Object.fromEntries(url.searchParams)).toMatchObject({
      roleId: '42', accountId: 'user', chatId: 'conversation', categoryId: '7', enableRag: 'true', rewriteStrategy: 'NONE',
    });
  });

  it('retains explicit rewrite configuration and disabled RAG', async () => {
    const config = parseConfig({ integrateId: '42', requestDomain: 'https://example.test', rewriteStrategy: 'MULTI_QUERY' })!;
    setApiConfig(config);
    vi.stubGlobal('fetch', vi.fn().mockImplementation(() => Promise.resolve(Response.json({ text: '', mcpEvents: [], suggestions: [], sources: [] }))));
    await chatRequest('问题', undefined, undefined, true);
    await chatRequest('问题', undefined, undefined, false);
    const urls = vi.mocked(fetch).mock.calls.map(([url]) => new URL(url as string));
    expect(urls[0]!.searchParams.get('rewriteStrategy')).toBe('MULTI_QUERY');
    expect(urls[1]!.searchParams.get('enableRag')).toBe('false');
  });
});