85 lines
3.8 KiB
JavaScript
85 lines
3.8 KiB
JavaScript
import { createServer } from 'node:http';
|
|
import { timingSafeEqual } from 'node:crypto';
|
|
import { chmod, rm } from 'node:fs/promises';
|
|
import { Readable } from 'node:stream';
|
|
|
|
export const MAX_REQUEST_BYTES = 32 * 1024 * 1024;
|
|
export function validRunToken(received, expected) {
|
|
const actual = Buffer.from(String(received || '').replace(/^Bearer /, ''));
|
|
const wanted = Buffer.from(expected);
|
|
return actual.length === wanted.length && timingSafeEqual(actual, wanted);
|
|
}
|
|
async function readJson(request) {
|
|
let size = 0; const chunks = [];
|
|
for await (const chunk of request) {
|
|
size += chunk.length;
|
|
if (size > MAX_REQUEST_BYTES) throw new Error('请求超出任务网关大小限制');
|
|
chunks.push(chunk);
|
|
}
|
|
return JSON.parse(Buffer.concat(chunks).toString('utf8'));
|
|
}
|
|
function json(response, status, value) {
|
|
response.writeHead(status, { 'content-type': 'application/json', 'cache-control': 'no-store' });
|
|
response.end(JSON.stringify(value));
|
|
}
|
|
|
|
// A broker belongs to one authenticated task. No user/project/model credentials
|
|
// or arbitrary upstream URL can be selected by its sandbox client.
|
|
export async function createRunBroker(options) {
|
|
const pending = new Set();
|
|
const server = createServer(async (request, response) => {
|
|
if (!validRunToken(request.headers.authorization, options.token)) {
|
|
json(response, 401, { error: '任务网关鉴权失败' }); return;
|
|
}
|
|
const controller = new AbortController(); pending.add(controller);
|
|
response.on('close', () => { if (!response.writableEnded) controller.abort(); });
|
|
try {
|
|
if (request.method !== 'POST') { json(response, 405, { error: '不支持的任务网关方法' }); return; }
|
|
const body = await readJson(request);
|
|
if (request.url === '/mcp') {
|
|
if (typeof body?.method !== 'string') throw new Error('无效 MCP 请求');
|
|
const result = await options.mcp(body, controller.signal);
|
|
json(response, 200, result ?? { notificationAccepted: true });
|
|
} else if (['/v1/chat/completions', '/v1/responses', '/v1/messages'].includes(request.url)) {
|
|
const upstream = await options.model(body, request.url, controller.signal);
|
|
response.writeHead(upstream.status, {
|
|
'content-type': upstream.headers.get('content-type') || 'application/json',
|
|
'cache-control': 'no-store',
|
|
});
|
|
if (upstream.body) await new Promise((resolve, reject) => {
|
|
const stream = Readable.fromWeb(upstream.body);
|
|
response.once('finish', resolve); response.once('close', resolve);
|
|
stream.once('error', reject); stream.pipe(response);
|
|
}); else response.end();
|
|
} else json(response, 404, { error: '任务网关入口不存在' });
|
|
} catch {
|
|
if (!response.headersSent) json(response, 502, { error: { message: '任务服务调用失败，请查看任务状态', type: 'canvas_agent_gateway_error' } });
|
|
else response.destroy();
|
|
} finally { pending.delete(controller); }
|
|
});
|
|
server.requestTimeout = 20 * 60_000;
|
|
server.headersTimeout = 30_000;
|
|
if (options.socketPath) {
|
|
await rm(options.socketPath, { force: true });
|
|
await new Promise((resolve, reject) => {
|
|
server.once('error', reject); server.listen(options.socketPath, resolve);
|
|
});
|
|
await chmod(options.socketPath, 0o600);
|
|
} else {
|
|
await new Promise((resolve, reject) => {
|
|
server.once('error', reject); server.listen(0, '127.0.0.1', resolve);
|
|
});
|
|
}
|
|
const address = server.address();
|
|
return {
|
|
url: typeof address === 'object' ? `http://127.0.0.1:${address.port}` : undefined,
|
|
socketPath: options.socketPath,
|
|
async close() {
|
|
for (const controller of pending) controller.abort();
|
|
server.closeAllConnections();
|
|
await new Promise(resolve => server.close(resolve));
|
|
if (options.socketPath) await rm(options.socketPath, { force: true });
|
|
},
|
|
};
|
|
}
|