483 lines
16 KiB
TypeScript
483 lines
16 KiB
TypeScript
import { EventEmitter } from 'node:events';
|
|
import type { IncomingMessage, ServerResponse } from 'http';
|
|
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
|
|
import { handleAiProxyRoutes } from '@electron/api/routes/ai-proxy';
|
|
import {
|
|
clearWorksSquareAIGatewayCredential,
|
|
seedWorksSquareAIGatewayCredential,
|
|
} from '@electron/services/works-square-ai-gateway';
|
|
import {
|
|
clearWorksSquareSession,
|
|
storeWorksSquareSession,
|
|
} from '@electron/services/works-square-session';
|
|
|
|
const loggerWarnMock = vi.hoisted(() => vi.fn());
|
|
|
|
vi.mock('@electron/utils/logger', () => ({
|
|
logger: {
|
|
debug: vi.fn(),
|
|
error: vi.fn(),
|
|
info: vi.fn(),
|
|
warn: (...args: unknown[]) => loggerWarnMock(...args),
|
|
},
|
|
}));
|
|
|
|
function createRequest(
|
|
method: string,
|
|
body?: unknown,
|
|
headers: Record<string, string> = {},
|
|
): IncomingMessage {
|
|
const req = new EventEmitter();
|
|
const rawBody = body === undefined
|
|
? undefined
|
|
: (typeof body === 'string' ? body : JSON.stringify(body));
|
|
Object.assign(req, {
|
|
method,
|
|
headers: {
|
|
host: '127.0.0.1:13210',
|
|
authorization: 'Bearer local-host-api-token',
|
|
'content-type': 'application/json',
|
|
...(rawBody ? { 'content-length': String(Buffer.byteLength(rawBody)) } : {}),
|
|
...headers,
|
|
},
|
|
[Symbol.asyncIterator]: async function* () {
|
|
if (rawBody !== undefined) {
|
|
yield Buffer.from(rawBody);
|
|
}
|
|
},
|
|
});
|
|
return req as IncomingMessage;
|
|
}
|
|
|
|
function createResponse(options: { writeResults?: boolean[] } = {}) {
|
|
const chunks: string[] = [];
|
|
const headers = new Map<string, number | string | string[]>();
|
|
const writeResults = [...(options.writeResults ?? [])];
|
|
const socket = { setNoDelay: vi.fn() };
|
|
const res = new EventEmitter();
|
|
Object.assign(res, {
|
|
statusCode: 0,
|
|
socket,
|
|
flushHeaders: vi.fn(),
|
|
setHeader: vi.fn((name: string, value: number | string | string[]) => {
|
|
headers.set(name.toLowerCase(), value);
|
|
}),
|
|
write: vi.fn((chunk: string | Uint8Array) => {
|
|
chunks.push(typeof chunk === 'string' ? chunk : Buffer.from(chunk).toString('utf8'));
|
|
return writeResults.shift() ?? true;
|
|
}),
|
|
end: vi.fn((chunk?: string | Uint8Array) => {
|
|
if (chunk !== undefined) {
|
|
chunks.push(typeof chunk === 'string' ? chunk : Buffer.from(chunk).toString('utf8'));
|
|
}
|
|
}),
|
|
});
|
|
|
|
return {
|
|
res: res as unknown as ServerResponse,
|
|
socket,
|
|
get statusCode() {
|
|
return (res as { statusCode: number }).statusCode;
|
|
},
|
|
header: (name: string) => headers.get(name.toLowerCase()),
|
|
body: () => chunks.join(''),
|
|
};
|
|
}
|
|
|
|
function streamFromText(...chunks: string[]): ReadableStream<Uint8Array> {
|
|
const encoder = new TextEncoder();
|
|
return new ReadableStream<Uint8Array>({
|
|
start(controller) {
|
|
for (const chunk of chunks) {
|
|
controller.enqueue(encoder.encode(chunk));
|
|
}
|
|
controller.close();
|
|
},
|
|
});
|
|
}
|
|
|
|
async function flushMicrotasks(): Promise<void> {
|
|
for (let index = 0; index < 20; index += 1) {
|
|
await Promise.resolve();
|
|
}
|
|
}
|
|
|
|
describe('ai proxy routes', () => {
|
|
beforeEach(() => {
|
|
vi.useFakeTimers();
|
|
vi.setSystemTime(new Date('2026-07-06T08:00:00.000Z'));
|
|
loggerWarnMock.mockReset();
|
|
clearWorksSquareAIGatewayCredential();
|
|
clearWorksSquareSession();
|
|
});
|
|
|
|
afterEach(() => {
|
|
vi.useRealTimers();
|
|
clearWorksSquareAIGatewayCredential();
|
|
clearWorksSquareSession();
|
|
vi.unstubAllGlobals();
|
|
});
|
|
|
|
it('forwards OpenAI-compatible JSON requests to one-api with a fresh Works Square gateway token', async () => {
|
|
seedWorksSquareAIGatewayCredential({
|
|
accessToken: 'ws-ai-token',
|
|
expiresIn: 3600,
|
|
oneApiBaseUrl: 'https://one-api.example.com/v1',
|
|
});
|
|
const fetchMock = vi.fn().mockResolvedValueOnce(
|
|
new Response(JSON.stringify({ id: 'chatcmpl_1' }), {
|
|
status: 200,
|
|
headers: { 'content-type': 'application/json' },
|
|
}),
|
|
);
|
|
vi.stubGlobal('fetch', fetchMock);
|
|
const response = createResponse();
|
|
|
|
const handled = await handleAiProxyRoutes(
|
|
createRequest('POST', { model: 'deepseek-chat', messages: [] }),
|
|
response.res,
|
|
new URL('http://127.0.0.1:13210/api/ai-proxy/v1/chat/completions?trace=1'),
|
|
{} as never,
|
|
);
|
|
|
|
expect(handled).toBe(true);
|
|
expect(fetchMock).toHaveBeenCalledWith(
|
|
'https://one-api.example.com/v1/chat/completions?trace=1',
|
|
expect.objectContaining({
|
|
method: 'POST',
|
|
body: Buffer.from(JSON.stringify({ model: 'deepseek-chat', messages: [] })),
|
|
headers: expect.objectContaining({
|
|
'content-type': 'application/json',
|
|
'x-works-square-ai-token': 'ws-ai-token',
|
|
}),
|
|
}),
|
|
);
|
|
const forwardedHeaders = fetchMock.mock.calls[0][1].headers as Record<string, string>;
|
|
expect(forwardedHeaders.authorization).toBeUndefined();
|
|
expect(forwardedHeaders.host).toBeUndefined();
|
|
expect(response.statusCode).toBe(200);
|
|
expect(response.header('content-type')).toBe('application/json');
|
|
expect(response.body()).toBe(JSON.stringify({ id: 'chatcmpl_1' }));
|
|
});
|
|
|
|
it('refreshes a near-expiry gateway credential before forwarding a long-running task request', async () => {
|
|
storeWorksSquareSession({
|
|
accessToken: 'works-access-token',
|
|
refreshToken: 'works-refresh-token',
|
|
expiresAt: Date.now() + 120_000,
|
|
});
|
|
seedWorksSquareAIGatewayCredential({
|
|
accessToken: 'near-expiry-ai-token',
|
|
expiresIn: 30,
|
|
oneApiBaseUrl: 'https://one-api.example.com/v1',
|
|
});
|
|
const fetchMock = vi.fn()
|
|
.mockResolvedValueOnce(new Response(JSON.stringify({
|
|
access_token: 'fresh-ai-token',
|
|
expires_in: 1800,
|
|
one_api_base_url: 'https://one-api.example.com/v1',
|
|
}), { status: 200 }))
|
|
.mockResolvedValueOnce(new Response(JSON.stringify({ id: 'chatcmpl_refresh' }), {
|
|
status: 200,
|
|
headers: { 'content-type': 'application/json' },
|
|
}));
|
|
vi.stubGlobal('fetch', fetchMock);
|
|
const response = createResponse();
|
|
|
|
await handleAiProxyRoutes(
|
|
createRequest('POST', { model: 'qwen3.7-max', messages: [] }),
|
|
response.res,
|
|
new URL('http://127.0.0.1:13210/api/ai-proxy/v1/chat/completions'),
|
|
{} as never,
|
|
);
|
|
|
|
expect(fetchMock).toHaveBeenCalledTimes(2);
|
|
expect(fetchMock.mock.calls[0][0]).toBe('https://square.nianxx.cn/api/ai-gateway/session');
|
|
expect(fetchMock.mock.calls[1][1].headers).toMatchObject({
|
|
'x-works-square-ai-token': 'fresh-ai-token',
|
|
});
|
|
expect(response.statusCode).toBe(200);
|
|
expect(response.body()).toBe(JSON.stringify({ id: 'chatcmpl_refresh' }));
|
|
});
|
|
|
|
it('refreshes the gateway credential and retries once when one-api rejects an expired gateway token', async () => {
|
|
storeWorksSquareSession({
|
|
accessToken: 'works-access-token',
|
|
refreshToken: 'works-refresh-token',
|
|
expiresAt: Date.now() + 120_000,
|
|
});
|
|
seedWorksSquareAIGatewayCredential({
|
|
accessToken: 'expired-ai-token',
|
|
expiresIn: 3600,
|
|
oneApiBaseUrl: 'https://one-api.example.com/v1',
|
|
});
|
|
const fetchMock = vi.fn()
|
|
.mockResolvedValueOnce(new Response('AI access token expired', { status: 401 }))
|
|
.mockResolvedValueOnce(
|
|
new Response(JSON.stringify({
|
|
access_token: 'fresh-ai-token',
|
|
expires_in: 1800,
|
|
one_api_base_url: 'https://one-api.example.com/v1',
|
|
}), { status: 200 }),
|
|
)
|
|
.mockResolvedValueOnce(new Response(JSON.stringify({ id: 'chatcmpl_2' }), { status: 200 }));
|
|
vi.stubGlobal('fetch', fetchMock);
|
|
const response = createResponse();
|
|
|
|
await handleAiProxyRoutes(
|
|
createRequest('POST', { model: 'deepseek-chat' }),
|
|
response.res,
|
|
new URL('http://127.0.0.1:13210/api/ai-proxy/v1/chat/completions'),
|
|
{} as never,
|
|
);
|
|
|
|
expect(fetchMock).toHaveBeenCalledTimes(3);
|
|
expect(fetchMock.mock.calls[0][1].headers).toMatchObject({
|
|
'x-works-square-ai-token': 'expired-ai-token',
|
|
});
|
|
expect(fetchMock.mock.calls[1][0]).toBe('https://square.nianxx.cn/api/ai-gateway/session');
|
|
expect(fetchMock.mock.calls[2][1].headers).toMatchObject({
|
|
'x-works-square-ai-token': 'fresh-ai-token',
|
|
});
|
|
expect(response.statusCode).toBe(200);
|
|
expect(response.body()).toBe(JSON.stringify({ id: 'chatcmpl_2' }));
|
|
});
|
|
|
|
it('does not retry quota failures', async () => {
|
|
seedWorksSquareAIGatewayCredential({
|
|
accessToken: 'ws-ai-token',
|
|
expiresIn: 3600,
|
|
oneApiBaseUrl: 'https://one-api.example.com/v1',
|
|
});
|
|
const fetchMock = vi.fn().mockResolvedValueOnce(
|
|
new Response(JSON.stringify({ error: 'user quota is not enough' }), {
|
|
status: 429,
|
|
headers: { 'content-type': 'application/json' },
|
|
}),
|
|
);
|
|
vi.stubGlobal('fetch', fetchMock);
|
|
const response = createResponse();
|
|
|
|
await handleAiProxyRoutes(
|
|
createRequest('POST', { model: 'deepseek-chat' }),
|
|
response.res,
|
|
new URL('http://127.0.0.1:13210/api/ai-proxy/v1/chat/completions'),
|
|
{} as never,
|
|
);
|
|
|
|
expect(fetchMock).toHaveBeenCalledOnce();
|
|
expect(response.statusCode).toBe(402);
|
|
expect(response.body()).toBe(JSON.stringify({ error: 'user quota is not enough' }));
|
|
});
|
|
|
|
it('maps rolling-window quota exhaustion to a non-retryable response status', async () => {
|
|
seedWorksSquareAIGatewayCredential({
|
|
accessToken: 'ws-ai-token',
|
|
expiresIn: 3600,
|
|
oneApiBaseUrl: 'https://one-api.example.com/v1',
|
|
});
|
|
const fetchMock = vi.fn().mockResolvedValueOnce(
|
|
new Response(JSON.stringify({
|
|
error: {
|
|
message: 'Token quota exhausted for rolling 5-hour window',
|
|
code: 'works_square_gateway_authorize_failed',
|
|
type: 'one_api_error',
|
|
},
|
|
}), {
|
|
status: 429,
|
|
headers: { 'content-type': 'application/json' },
|
|
}),
|
|
);
|
|
vi.stubGlobal('fetch', fetchMock);
|
|
const response = createResponse();
|
|
|
|
await handleAiProxyRoutes(
|
|
createRequest('POST', { model: 'qwen3.6-plus', messages: [] }),
|
|
response.res,
|
|
new URL('http://127.0.0.1:13210/api/ai-proxy/v1/chat/completions'),
|
|
{} as never,
|
|
);
|
|
|
|
expect(fetchMock).toHaveBeenCalledOnce();
|
|
expect(response.statusCode).toBe(402);
|
|
expect(response.body()).toContain('Token quota exhausted for rolling 5-hour window');
|
|
});
|
|
|
|
it('preserves the upstream status for a generic gateway authorization failure', async () => {
|
|
seedWorksSquareAIGatewayCredential({
|
|
accessToken: 'ws-ai-token',
|
|
expiresIn: 3600,
|
|
oneApiBaseUrl: 'https://one-api.example.com/v1',
|
|
});
|
|
const fetchMock = vi.fn().mockResolvedValueOnce(
|
|
new Response(JSON.stringify({
|
|
error: {
|
|
message: 'works_square_gateway_authorize_failed',
|
|
code: 'works_square_gateway_authorize_failed',
|
|
type: 'one_api_error',
|
|
},
|
|
}), {
|
|
status: 429,
|
|
headers: { 'content-type': 'application/json' },
|
|
}),
|
|
);
|
|
vi.stubGlobal('fetch', fetchMock);
|
|
const response = createResponse();
|
|
|
|
await handleAiProxyRoutes(
|
|
createRequest('POST', { model: 'qwen3.6-plus', messages: [] }),
|
|
response.res,
|
|
new URL('http://127.0.0.1:13210/api/ai-proxy/v1/chat/completions'),
|
|
{} as never,
|
|
);
|
|
|
|
expect(fetchMock).toHaveBeenCalledOnce();
|
|
expect(response.statusCode).toBe(429);
|
|
expect(response.body()).toContain('works_square_gateway_authorize_failed');
|
|
});
|
|
|
|
it('logs a sanitized one-api error summary when the upstream group is saturated', async () => {
|
|
seedWorksSquareAIGatewayCredential({
|
|
accessToken: 'ws-ai-token',
|
|
expiresIn: 3600,
|
|
oneApiBaseUrl: 'https://one-api.example.com/v1',
|
|
});
|
|
const fetchMock = vi.fn().mockResolvedValueOnce(
|
|
new Response(JSON.stringify({
|
|
error: {
|
|
message: '当前分组上游负载已饱和，请稍后再试 (request id: req-saturated)',
|
|
code: 'rate_limit_exceeded',
|
|
type: 'one_api_error',
|
|
},
|
|
}), {
|
|
status: 429,
|
|
headers: { 'content-type': 'application/json' },
|
|
}),
|
|
);
|
|
vi.stubGlobal('fetch', fetchMock);
|
|
const response = createResponse();
|
|
|
|
await handleAiProxyRoutes(
|
|
createRequest('POST', { model: 'qwen3.7-max', messages: [{ role: 'user', content: 'secret prompt' }] }),
|
|
response.res,
|
|
new URL('http://127.0.0.1:13210/api/ai-proxy/v1/chat/completions'),
|
|
{} as never,
|
|
);
|
|
|
|
expect(response.statusCode).toBe(429);
|
|
expect(response.body()).toContain('当前分组上游负载已饱和');
|
|
expect(loggerWarnMock).toHaveBeenCalledWith(
|
|
'[ai-proxy] One-api returned non-success response',
|
|
expect.objectContaining({
|
|
status: 429,
|
|
url: 'https://one-api.example.com/v1/chat/completions',
|
|
code: 'rate_limit_exceeded',
|
|
type: 'one_api_error',
|
|
message: '当前分组上游负载已饱和，请稍后再试 (request id: req-saturated)',
|
|
}),
|
|
);
|
|
expect(JSON.stringify(loggerWarnMock.mock.calls)).not.toContain('secret prompt');
|
|
expect(JSON.stringify(loggerWarnMock.mock.calls)).not.toContain('ws-ai-token');
|
|
});
|
|
|
|
it('strips decoded compression headers from proxied responses', async () => {
|
|
seedWorksSquareAIGatewayCredential({
|
|
accessToken: 'ws-ai-token',
|
|
expiresIn: 3600,
|
|
oneApiBaseUrl: 'https://one-api.example.com/v1',
|
|
});
|
|
const fetchMock = vi.fn().mockResolvedValueOnce(
|
|
new Response(JSON.stringify({ id: 'chatcmpl_br' }), {
|
|
status: 200,
|
|
headers: {
|
|
'content-encoding': 'br',
|
|
'content-length': '999',
|
|
'content-type': 'application/json',
|
|
},
|
|
}),
|
|
);
|
|
vi.stubGlobal('fetch', fetchMock);
|
|
const response = createResponse();
|
|
|
|
await handleAiProxyRoutes(
|
|
createRequest('POST', { model: 'deepseek-chat' }, { 'accept-encoding': 'br, gzip' }),
|
|
response.res,
|
|
new URL('http://127.0.0.1:13210/api/ai-proxy/v1/chat/completions'),
|
|
{} as never,
|
|
);
|
|
|
|
const forwardedHeaders = fetchMock.mock.calls[0][1].headers as Record<string, string>;
|
|
expect(forwardedHeaders['accept-encoding']).toBeUndefined();
|
|
expect(response.statusCode).toBe(200);
|
|
expect(response.header('content-type')).toBe('application/json');
|
|
expect(response.header('content-encoding')).toBeUndefined();
|
|
expect(response.header('content-length')).toBeUndefined();
|
|
expect(response.body()).toBe(JSON.stringify({ id: 'chatcmpl_br' }));
|
|
});
|
|
|
|
it('passes streaming chunks through to the local client', async () => {
|
|
seedWorksSquareAIGatewayCredential({
|
|
accessToken: 'ws-ai-token',
|
|
expiresIn: 3600,
|
|
oneApiBaseUrl: 'https://one-api.example.com/v1',
|
|
});
|
|
const fetchMock = vi.fn().mockResolvedValueOnce(
|
|
new Response(streamFromText('data: one\n\n', 'data: two\n\n'), {
|
|
status: 200,
|
|
headers: { 'content-type': 'text/event-stream' },
|
|
}),
|
|
);
|
|
vi.stubGlobal('fetch', fetchMock);
|
|
const response = createResponse();
|
|
|
|
await handleAiProxyRoutes(
|
|
createRequest('POST', { stream: true }),
|
|
response.res,
|
|
new URL('http://127.0.0.1:13210/api/ai-proxy/v1/chat/completions'),
|
|
{} as never,
|
|
);
|
|
|
|
expect(response.statusCode).toBe(200);
|
|
expect(response.header('content-type')).toBe('text/event-stream');
|
|
expect(response.body()).toBe('data: one\n\ndata: two\n\n');
|
|
});
|
|
|
|
it('flushes streaming headers and waits for downstream drain before reading more chunks', async () => {
|
|
seedWorksSquareAIGatewayCredential({
|
|
accessToken: 'ws-ai-token',
|
|
expiresIn: 3600,
|
|
oneApiBaseUrl: 'https://one-api.example.com/v1',
|
|
});
|
|
const fetchMock = vi.fn().mockResolvedValueOnce(
|
|
new Response(streamFromText('data: one\n\n', 'data: two\n\n'), {
|
|
status: 200,
|
|
headers: { 'content-type': 'text/event-stream' },
|
|
}),
|
|
);
|
|
vi.stubGlobal('fetch', fetchMock);
|
|
const response = createResponse({ writeResults: [false, true] });
|
|
let settled = false;
|
|
|
|
const handledPromise = handleAiProxyRoutes(
|
|
createRequest('POST', { stream: true }),
|
|
response.res,
|
|
new URL('http://127.0.0.1:13210/api/ai-proxy/v1/chat/completions'),
|
|
{} as never,
|
|
).finally(() => {
|
|
settled = true;
|
|
});
|
|
await flushMicrotasks();
|
|
|
|
expect(response.res.flushHeaders).toHaveBeenCalledOnce();
|
|
expect(response.socket.setNoDelay).toHaveBeenCalledWith(true);
|
|
expect(response.body()).toBe('data: one\n\n');
|
|
expect(settled).toBe(false);
|
|
|
|
response.res.emit('drain');
|
|
await handledPromise;
|
|
|
|
expect(settled).toBe(true);
|
|
expect(response.body()).toBe('data: one\n\ndata: two\n\n');
|
|
});
|
|
});
|