zsjz-ai/ai-electron/electron/service/codex/chatRouter/routerService.test.ts

import { createServer, type Server } from 'node:http';
import { afterAll, beforeAll, afterEach, describe, expect, it } from 'vitest';
import { ChatRouterService } from './routerService';

let lastUpstreamBody: string | null = null;
let lastUpstreamHeaders: Record<string, string | string[] | undefined> = {};
let upstreamAbortCount = 0;
/** 上游行为：'sse' 正常流 | 'slow' 慢 prefill | 'error' 500 */
let upstreamMode = 'sse';

const upstream: Server = createServer((req, res) => {
  const chunks: Buffer[] = [];
  req.on('data', (c: Buffer) => chunks.push(c));
  req.on('end', () => {
    lastUpstreamBody = Buffer.concat(chunks).toString('utf8');
    lastUpstreamHeaders = req.headers;
    if (upstreamMode === 'error') {
      res.writeHead(500, { 'content-type': 'application/json' });
      res.end(JSON.stringify({ error: { message: 'boom' } }));
      return;
    }
    // 非流式请求（防御路径）：回 JSON chat completion；流式回 SSE
    const wantsStream = (() => {
      try { return (JSON.parse(lastUpstreamBody ?? '{}') as { stream?: boolean }).stream === true; } catch { return true; }
    })();
    if (!wantsStream) {
      res.writeHead(200, { 'content-type': 'application/json' });
      res.end(JSON.stringify({
        choices: [{ message: { content: '你好世界' }, finish_reason: 'stop' }],
        usage: { prompt_tokens: 5, completion_tokens: 2, total_tokens: 7 },
      }));
      return;
    }
    res.writeHead(200, { 'content-type': 'text/event-stream' });
    const finish = (): void => {
      res.write(`data: ${JSON.stringify({ choices: [{ delta: { content: '世界' } }, { finish_reason: 'stop' }] })}\n\n`);
      res.write(`data: ${JSON.stringify({ choices: [], usage: { prompt_tokens: 5, completion_tokens: 2, total_tokens: 7 } })}\n\n`);
      res.write('data: [DONE]\n\n');
      res.end();
    };
    if (upstreamMode === 'slow') {
      // 慢 prefill：3 秒后才发首字节，路由层必须在这之前就把 SSE 头与 begin 事件发出去
      setTimeout(() => {
        res.write(`data: ${JSON.stringify({ choices: [{ delta: { content: '晚到' } }] })}\n\n`);
        finish();
      }, 3_000);
    } else {
      res.write(`data: ${JSON.stringify({ choices: [{ delta: { content: '你好' } }] })}\n\n`);
      finish();
    }
  });
  req.on('close', () => {
    if (!res.writableEnded) upstreamAbortCount += 1;
  });
});

function parseSse(raw: string): Array<{ event: string; data: Record<string, unknown> }> {
  return raw.split('\n\n').filter((part) => part.includes('event:')).map((part) => {
    const event = /^event: (.*)$/mu.exec(part)?.[1] ?? '';
    const data = /^data: (.*)$/mu.exec(part)?.[1] ?? '{}';
    return { event, data: JSON.parse(data) as Record<string, unknown> };
  });
}

let upstreamUrl = '';
const routers: ChatRouterService[] = [];

beforeAll(async () => {
  await new Promise<void>((resolve) => {
    upstream.listen(0, '127.0.0.1', () => resolve());
  });
  const address = upstream.address();
  upstreamUrl = `http://127.0.0.1:${(address as { port: number }).port}/v1`;
});

async function startRouter(mode: string): Promise<ChatRouterService> {
  upstreamMode = mode;
  const router = new ChatRouterService();
  routers.push(router);
  await router.start({
    upstreamBaseUrl: upstreamUrl,
    forwardHeaderNames: ['x-extra'],
    onDiagnostic: () => {},
  });
  return router;
}

afterEach(async () => {
  upstreamAbortCount = 0;
  for (const router of routers.splice(0)) await router.stop();
});

afterAll(() => {
  upstream.close();
});

describe('ChatRouterService', () => {
  it('非流式：上游收到 /chat/completions 与透传头，下游拿到 Responses JSON', async () => {
    const router = await startRouter('sse');
    // router.url 本身以 /v1 结尾（codex 语义），所以请求路径拼 /responses
    const res = await fetch(`${router.url}/responses`, {
      method: 'POST',
      headers: { 'content-type': 'application/json', authorization: 'Bearer sk-test', 'x-extra': 'v1' },
      body: JSON.stringify({ model: 'qwen', instructions: '你是助手', input: '你好', stream: false }),
    });
    expect(res.status).toBe(200);
    const payload = await res.json() as Record<string, unknown>;
    expect(payload.object).toBe('response');
    expect(payload.status).toBe('completed');
    expect((payload.usage as Record<string, unknown>).total_tokens).toBe(7);

    expect(lastUpstreamBody).not.toBeNull();
    const chatRequest = JSON.parse(lastUpstreamBody!) as Record<string, unknown>;
    expect(chatRequest.stream).toBe(false);
    expect(chatRequest.stream_options).toBeUndefined();
    expect(chatRequest.messages).toEqual([
      { role: 'developer', content: '你是助手' },
      { role: 'user', content: '你好' },
    ]);
    expect(lastUpstreamHeaders.authorization).toBe('Bearer sk-test');
    expect(lastUpstreamHeaders['x-extra']).toBe('v1');
  });

  it('流式：先 created 后 completed，delta 拼接正确，上游 body 带 include_usage', async () => {
    const router = await startRouter('sse');
    const res = await fetch(`${router.url}/responses`, {
      method: 'POST',
      headers: { 'content-type': 'application/json', authorization: 'Bearer sk-stream' },
      body: JSON.stringify({ model: 'qwen', input: '你好', stream: true }),
    });
    expect(res.status).toBe(200);
    expect(res.headers.get('content-type')).toContain('text/event-stream');
    const raw = await res.text();
    const events = parseSse(raw);
    expect(events[0].event).toBe('response.created');
    expect(events.at(-1)?.event).toBe('response.completed');
    const text = events.filter((e) => e.event === 'response.output_text.delta')
      .map((e) => e.data.delta as string).join('');
    expect(text).toBe('你好世界');
    const completed = events.at(-1)!.data.response as Record<string, unknown>;
    expect(completed.id).toMatch(/^resp_/);
    expect(completed.usage).toMatchObject({ input_tokens: 5, output_tokens: 2, total_tokens: 7 });

    expect(lastUpstreamBody).not.toBeNull();
    const chatRequest = JSON.parse(lastUpstreamBody!) as Record<string, unknown>;
    expect(chatRequest.stream).toBe(true);
    expect(chatRequest.stream_options).toEqual({ include_usage: true });
    expect(lastUpstreamHeaders.authorization).toBe('Bearer sk-stream');
  });

  it('上游 500（流式）：以 response.failed 收尾，codex 走错误映射', async () => {
    const router = await startRouter('error');
    const res = await fetch(`${router.url}/responses`, {
      method: 'POST',
      headers: { 'content-type': 'application/json' },
      body: JSON.stringify({ model: 'qwen', input: 'hi', stream: true }),
    });
    expect(res.status).toBe(200);
    const events = parseSse(await res.text());
    expect(events.at(-1)?.event).toBe('response.failed');
    expect((events.at(-1)!.data.response as Record<string, unknown>).error).toMatchObject({ code: 'upstream_error' });
  });

  it('只接受 POST /v1/responses；stop 后端口释放', async () => {
    const router = await startRouter('sse');
    const getUrl = `${router.url}/models`;
    const wrong = await fetch(getUrl);
    expect(wrong.status).toBe(404);
    const url = router.url;
    await router.stop();
    await expect(fetch(url!)).rejects.toThrow();
  });

  it('客户端提前挂断：1 秒内取消上游请求', async () => {
    const router = await startRouter('slow');
    const controller = new AbortController();
    const res = await fetch(`${router.url}/responses`, {
      method: 'POST',
      headers: { 'content-type': 'application/json' },
      body: JSON.stringify({ model: 'qwen', input: 'hi', stream: true }),
      signal: controller.signal,
    });
    await res.arrayBuffer().catch(() => undefined);
    controller.abort();
    await new Promise((resolve) => setTimeout(resolve, 1_000));
    expect(upstreamAbortCount).toBeGreaterThanOrEqual(1);
  }, 10_000);

  it('慢 prefill：上游 3 秒不响应，下游 1.5 秒内已收到 begin 事件不掐线', async () => {
    const router = await startRouter('slow');
    const startedAt = Date.now();
    const res = await fetch(`${router.url}/responses`, {
      method: 'POST',
      headers: { 'content-type': 'application/json' },
      body: JSON.stringify({ model: 'qwen', input: 'hi', stream: true }),
    });
    const reader = res.body!.getReader();
    const decoder = new TextDecoder();
    let firstFrameAt = -1;
    let raw = '';
    for (;;) {
      const { done, value } = await reader.read();
      if (done) break;
      raw += decoder.decode(value, { stream: true });
      if (firstFrameAt < 0 && raw.includes('response.created')) {
        firstFrameAt = Date.now() - startedAt;
        break; // 拿到 begin 即验证完毕，不等慢 prefill
      }
    }
    expect(firstFrameAt).toBeGreaterThanOrEqual(0);
    expect(firstFrameAt).toBeLessThan(1_500);
    reader.cancel().catch(() => undefined);
  }, 15_000);
});