feat: add streaming runtime router
This commit is contained in:
John
+263
-23
@@ -1,6 +1,7 @@
|
||||
import fs from 'node:fs';
|
||||
import path from 'node:path';
|
||||
import { Readable, Transform } from 'node:stream';
|
||||
import { Readable, Transform, Writable } from 'node:stream';
|
||||
import { pipeline } from 'node:stream/promises';
|
||||
import { Agent, fetch as undiciFetch } from 'undici';
|
||||
import { appendBalanceEvent, createSseBillingTransform } from './sse-billing.mjs';
|
||||
import { developerToolsFromPolicy } from './capabilities.mjs';
|
||||
@@ -79,6 +80,164 @@ async function apiFetch(target, apiSecret, pathname, init = {}) {
|
||||
});
|
||||
}
|
||||
|
||||
function createRuntimeRouter({
|
||||
targets,
|
||||
targetHealthy,
|
||||
redisUrl = process.env.MEMIND_RUNTIME_REDIS_URL,
|
||||
namespace = process.env.MEMIND_RUNTIME_REDIS_NAMESPACE ?? 'memind:runtime',
|
||||
}) {
|
||||
if (!redisUrl || targets.length <= 1) return null;
|
||||
let clientPromise = null;
|
||||
let warned = false;
|
||||
const workerIdForTarget = (target) => {
|
||||
const idx = targets.indexOf(target);
|
||||
return idx >= 0 ? `goosed-${idx + 1}` : `goosed-${Buffer.from(String(target)).toString('base64url')}`;
|
||||
};
|
||||
const key = (...parts) => [namespace, ...parts].join(':');
|
||||
const getClient = async () => {
|
||||
if (!clientPromise) {
|
||||
clientPromise = import('redis')
|
||||
.then(async ({ createClient }) => {
|
||||
const client = createClient({ url: redisUrl });
|
||||
client.on('error', (err) => {
|
||||
if (!warned) {
|
||||
warned = true;
|
||||
console.warn('[RuntimeRouter] Redis error:', err instanceof Error ? err.message : err);
|
||||
}
|
||||
});
|
||||
await client.connect();
|
||||
console.log('[RuntimeRouter] Redis scheduler enabled');
|
||||
return client;
|
||||
})
|
||||
.catch((err) => {
|
||||
clientPromise = null;
|
||||
if (!warned) {
|
||||
warned = true;
|
||||
console.warn('[RuntimeRouter] Redis disabled:', err instanceof Error ? err.message : err);
|
||||
}
|
||||
return null;
|
||||
});
|
||||
}
|
||||
return clientPromise;
|
||||
};
|
||||
const readNumber = (value) => {
|
||||
const n = Number(value ?? 0);
|
||||
return Number.isFinite(n) ? n : 0;
|
||||
};
|
||||
const scoreWorker = async (client, target) => {
|
||||
const workerId = workerIdForTarget(target);
|
||||
const values = await client.mGet([
|
||||
key('worker', workerId, 'active_streams'),
|
||||
key('worker', workerId, 'active_sessions'),
|
||||
key('worker', workerId, 'ewma_first_token_ms'),
|
||||
key('worker', workerId, 'error_rate'),
|
||||
key('worker', workerId, 'memory_pressure'),
|
||||
key('worker', workerId, 'drain'),
|
||||
]);
|
||||
if (/^(1|true|yes)$/i.test(String(values[5] ?? ''))) return Number.POSITIVE_INFINITY;
|
||||
return (
|
||||
readNumber(values[0]) * 3 +
|
||||
readNumber(values[1]) +
|
||||
readNumber(values[2]) * 0.01 +
|
||||
readNumber(values[3]) * 5 +
|
||||
readNumber(values[4]) * 2
|
||||
);
|
||||
};
|
||||
return {
|
||||
async pickTarget(fallbackPick, orderedTargets = targets) {
|
||||
const client = await getClient();
|
||||
if (!client) return fallbackPick();
|
||||
let winner = null;
|
||||
for (const target of orderedTargets) {
|
||||
if (!await targetHealthy(target)) continue;
|
||||
const score = await scoreWorker(client, target).catch(() => Number.POSITIVE_INFINITY);
|
||||
if (!Number.isFinite(score)) continue;
|
||||
if (!winner || score < winner.score) winner = { target, score };
|
||||
}
|
||||
return winner?.target ?? fallbackPick();
|
||||
},
|
||||
async registerSession(sessionId, target) {
|
||||
const client = await getClient();
|
||||
if (!client || !sessionId || !target) return;
|
||||
const workerId = workerIdForTarget(target);
|
||||
await client
|
||||
.multi()
|
||||
.set(key('session', sessionId, 'worker'), workerId, { EX: 60 * 60 * 24 * 30 })
|
||||
.set(key('session', sessionId, 'target'), target, { EX: 60 * 60 * 24 * 30 })
|
||||
.exec()
|
||||
.catch(() => null);
|
||||
},
|
||||
async resolveSessionTarget(sessionId) {
|
||||
const client = await getClient();
|
||||
if (!client || !sessionId) return null;
|
||||
const target = await client.get(key('session', sessionId, 'target')).catch(() => null);
|
||||
return target && targets.includes(target) ? target : null;
|
||||
},
|
||||
async streamStarted(sessionId, target) {
|
||||
const client = await getClient();
|
||||
if (!client || !target) return;
|
||||
const workerId = workerIdForTarget(target);
|
||||
const streamKey = sessionId ? key('stream', sessionId, 'status') : null;
|
||||
const multi = client
|
||||
.multi()
|
||||
.incr(key('worker', workerId, 'active_streams'))
|
||||
.set(key('worker', workerId, 'heartbeat'), String(Date.now()), { EX: 30 });
|
||||
if (streamKey) multi.set(streamKey, 'active', { EX: 60 * 60 });
|
||||
await multi.exec().catch(() => null);
|
||||
},
|
||||
async streamEnded(sessionId, target) {
|
||||
const client = await getClient();
|
||||
if (!client || !target) return;
|
||||
const workerId = workerIdForTarget(target);
|
||||
const streamKey = sessionId ? key('stream', sessionId, 'status') : null;
|
||||
const activeKey = key('worker', workerId, 'active_streams');
|
||||
const nextValue = await client.decr(activeKey).catch(() => null);
|
||||
const multi = client
|
||||
.multi()
|
||||
.set(key('worker', workerId, 'heartbeat'), String(Date.now()), { EX: 30 });
|
||||
if (Number(nextValue ?? 0) < 0) multi.set(activeKey, '0');
|
||||
if (streamKey) multi.set(streamKey, 'closed', { EX: 600 });
|
||||
await multi.exec().catch(() => null);
|
||||
},
|
||||
async getStatus() {
|
||||
const client = await getClient();
|
||||
const workers = [];
|
||||
for (const target of targets) {
|
||||
const workerId = workerIdForTarget(target);
|
||||
const values = client
|
||||
? await client
|
||||
.mGet([
|
||||
key('worker', workerId, 'active_streams'),
|
||||
key('worker', workerId, 'active_sessions'),
|
||||
key('worker', workerId, 'ewma_first_token_ms'),
|
||||
key('worker', workerId, 'error_rate'),
|
||||
key('worker', workerId, 'memory_pressure'),
|
||||
key('worker', workerId, 'heartbeat'),
|
||||
key('worker', workerId, 'drain'),
|
||||
])
|
||||
.catch(() => [])
|
||||
: [];
|
||||
workers.push({
|
||||
id: workerId,
|
||||
target,
|
||||
activeStreams: readNumber(values?.[0]),
|
||||
activeSessions: readNumber(values?.[1]),
|
||||
ewmaFirstTokenMs: readNumber(values?.[2]),
|
||||
errorRate: readNumber(values?.[3]),
|
||||
memoryPressure: readNumber(values?.[4]),
|
||||
heartbeat: values?.[5] ? Number(values[5]) : null,
|
||||
drain: /^(1|true|yes)$/i.test(String(values?.[6] ?? '')),
|
||||
});
|
||||
}
|
||||
return {
|
||||
enabled: Boolean(client),
|
||||
namespace,
|
||||
workers,
|
||||
};
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
async function readJsonBody(req) {
|
||||
if (req.method === 'GET' || req.method === 'HEAD') return null;
|
||||
const chunks = [];
|
||||
@@ -566,6 +725,7 @@ export function createTkmindProxy({
|
||||
return false;
|
||||
}
|
||||
}
|
||||
const runtimeRouter = createRuntimeRouter({ targets, targetHealthy });
|
||||
|
||||
function isGeneratedSessionName(name) {
|
||||
const normalized = typeof name === 'string' ? name.trim() : '';
|
||||
@@ -675,12 +835,56 @@ export function createTkmindProxy({
|
||||
|
||||
async function pickTarget() {
|
||||
if (targets.length <= 1) return primaryTarget;
|
||||
for (let i = 0; i < targets.length; i += 1) {
|
||||
const target = targets[rrIdx];
|
||||
const fallbackPick = async () => {
|
||||
for (let i = 0; i < targets.length; i += 1) {
|
||||
const target = targets[rrIdx];
|
||||
rrIdx = (rrIdx + 1) % targets.length;
|
||||
if (await targetHealthy(target)) return target;
|
||||
}
|
||||
return primaryTarget;
|
||||
};
|
||||
if (runtimeRouter) {
|
||||
const orderedTargets = targets.map((_, i) => targets[(rrIdx + i) % targets.length]);
|
||||
rrIdx = (rrIdx + 1) % targets.length;
|
||||
if (await targetHealthy(target)) return target;
|
||||
return runtimeRouter.pickTarget(fallbackPick, orderedTargets);
|
||||
}
|
||||
return primaryTarget;
|
||||
return fallbackPick();
|
||||
}
|
||||
|
||||
async function rememberSessionTarget(sessionId, target) {
|
||||
if (!runtimeRouter || !sessionId || !target) return;
|
||||
await runtimeRouter.registerSession(sessionId, target).catch(() => null);
|
||||
}
|
||||
|
||||
async function markStreamStarted(sessionId, target) {
|
||||
if (!runtimeRouter || !sessionId || !target) return;
|
||||
await runtimeRouter.streamStarted(sessionId, target).catch(() => null);
|
||||
}
|
||||
|
||||
async function markStreamEnded(sessionId, target) {
|
||||
if (!runtimeRouter || !sessionId || !target) return;
|
||||
await runtimeRouter.streamEnded(sessionId, target).catch(() => null);
|
||||
}
|
||||
|
||||
async function getRuntimeStatus() {
|
||||
const targetStatuses = [];
|
||||
for (const target of targets) {
|
||||
targetStatuses.push({
|
||||
target,
|
||||
healthy: await targetHealthy(target),
|
||||
});
|
||||
}
|
||||
const routerStatus = runtimeRouter
|
||||
? await runtimeRouter.getStatus().catch((err) => ({
|
||||
enabled: false,
|
||||
error: err instanceof Error ? err.message : String(err),
|
||||
}))
|
||||
: { enabled: false, reason: 'MEMIND_RUNTIME_REDIS_URL not configured' };
|
||||
return {
|
||||
publicBaseUrl: process.env.H5_PUBLIC_BASE_URL ?? null,
|
||||
router: routerStatus,
|
||||
targets: targetStatuses,
|
||||
};
|
||||
}
|
||||
|
||||
async function startSessionForUser(
|
||||
@@ -713,6 +917,7 @@ export function createTkmindProxy({
|
||||
if (!session?.id) {
|
||||
throw new Error('创建会话失败：缺少 session id');
|
||||
}
|
||||
await rememberSessionTarget(session.id, startTarget);
|
||||
await userAuth.registerAgentSession(userId, session.id, startTarget);
|
||||
if (resolvedSessionPolicy?.gooseMode) {
|
||||
const modeRes = await apiFetch(startTarget, apiSecret, '/agent/update_session', {
|
||||
@@ -756,6 +961,8 @@ export function createTkmindProxy({
|
||||
async function resolveTarget(sessionId) {
|
||||
if (targets.length <= 1 || !sessionId) return primaryTarget;
|
||||
try {
|
||||
const routedTarget = await runtimeRouter?.resolveSessionTarget(sessionId);
|
||||
if (routedTarget) return routedTarget;
|
||||
const { target, node } = await userAuth.getSessionTarget(sessionId);
|
||||
// Prefer the pinned URL: stable across reordering/resizing the target list.
|
||||
// Only honor it if that upstream is still configured; otherwise fall back to
|
||||
@@ -975,6 +1182,7 @@ export function createTkmindProxy({
|
||||
}
|
||||
const session = JSON.parse(text);
|
||||
if (session?.id) {
|
||||
await rememberSessionTarget(session.id, startTarget);
|
||||
await userAuth.registerAgentSession(
|
||||
req.currentUser.id,
|
||||
session.id,
|
||||
@@ -1228,16 +1436,25 @@ export function createTkmindProxy({
|
||||
];
|
||||
|
||||
const proxySessionEvents = async (req, res, sessionId, { onAfterFinish, onEvent } = {}) => {
|
||||
const upstreamAbort = new AbortController();
|
||||
let clientClosed = false;
|
||||
const abortUpstream = () => {
|
||||
clientClosed = true;
|
||||
upstreamAbort.abort();
|
||||
};
|
||||
req.once('close', abortUpstream);
|
||||
try {
|
||||
const pathname = `/sessions/${sessionId}/events`;
|
||||
const sessionTarget = await resolveTarget(sessionId);
|
||||
const upstream = await apiFetch(sessionTarget, apiSecret, pathname, {
|
||||
method: 'GET',
|
||||
signal: upstreamAbort.signal,
|
||||
headers: {
|
||||
Accept: 'text/event-stream',
|
||||
'Last-Event-ID': req.get('last-event-id') ?? '',
|
||||
},
|
||||
});
|
||||
if (clientClosed) return;
|
||||
|
||||
if (!upstream.ok || !upstream.body) {
|
||||
const text = await upstream.text().catch(() => '');
|
||||
@@ -1246,9 +1463,11 @@ export function createTkmindProxy({
|
||||
}
|
||||
|
||||
res.status(upstream.status);
|
||||
res.setHeader('Content-Type', 'text/event-stream');
|
||||
res.setHeader('Cache-Control', 'no-cache');
|
||||
res.setHeader('Content-Type', 'text/event-stream; charset=utf-8');
|
||||
res.setHeader('Cache-Control', 'no-cache, no-transform');
|
||||
res.setHeader('Connection', 'keep-alive');
|
||||
res.setHeader('X-Accel-Buffering', 'no');
|
||||
res.flushHeaders?.();
|
||||
|
||||
let pendingBalance = null;
|
||||
const billingTransform = createSseBillingTransform({
|
||||
@@ -1280,27 +1499,47 @@ export function createTkmindProxy({
|
||||
|
||||
const source = Readable.fromWeb(upstream.body);
|
||||
const linkSanitizer = createSessionEventSanitizer(req.currentUser, { onEvent });
|
||||
const waitForDrain = () => new Promise((resolve) => res.once('drain', resolve));
|
||||
const writeClientChunk = async (chunk) => {
|
||||
if (res.writableEnded || clientClosed) return;
|
||||
let needsDrain = !res.write(chunk);
|
||||
if (pendingBalance != null && !res.writableEnded && !clientClosed) {
|
||||
needsDrain = !res.write(appendBalanceEvent(pendingBalance)) || needsDrain;
|
||||
pendingBalance = null;
|
||||
}
|
||||
if (needsDrain && !res.writableEnded && !clientClosed) {
|
||||
await waitForDrain();
|
||||
}
|
||||
};
|
||||
const clientSink = new Writable({
|
||||
write(chunk, _encoding, callback) {
|
||||
writeClientChunk(chunk).then(() => callback(), callback);
|
||||
},
|
||||
});
|
||||
|
||||
// 每 20s 发一个注释行保持连接，防止 nginx/代理因静默超时切断 SSE
|
||||
const keepalive = setInterval(() => {
|
||||
if (!res.writableEnded) res.write(': keepalive\n\n');
|
||||
}, 20000);
|
||||
const stopKeepalive = () => clearInterval(keepalive);
|
||||
|
||||
billingTransform.on('data', (chunk) => {
|
||||
res.write(chunk);
|
||||
if (pendingBalance != null) {
|
||||
res.write(appendBalanceEvent(pendingBalance));
|
||||
pendingBalance = null;
|
||||
if (!res.writableEnded && !clientClosed) {
|
||||
const ok = res.write(': keepalive\n\n');
|
||||
if (!ok) source.pause();
|
||||
}
|
||||
});
|
||||
billingTransform.on('end', () => { stopKeepalive(); res.end(); });
|
||||
billingTransform.on('error', () => { stopKeepalive(); res.end(); });
|
||||
linkSanitizer.on('error', () => { stopKeepalive(); res.end(); });
|
||||
source.on('error', () => { stopKeepalive(); res.end(); });
|
||||
req.on('close', stopKeepalive);
|
||||
source.pipe(linkSanitizer).pipe(billingTransform);
|
||||
}, 20000);
|
||||
res.on('drain', () => source.resume());
|
||||
await markStreamStarted(sessionId, sessionTarget);
|
||||
try {
|
||||
await pipeline(source, linkSanitizer, billingTransform, clientSink);
|
||||
} finally {
|
||||
clearInterval(keepalive);
|
||||
req.off('close', abortUpstream);
|
||||
await markStreamEnded(sessionId, sessionTarget);
|
||||
if (!res.writableEnded) res.end();
|
||||
}
|
||||
} catch (err) {
|
||||
req.off('close', abortUpstream);
|
||||
if (clientClosed || upstreamAbort.signal.aborted) {
|
||||
if (!res.writableEnded) res.end();
|
||||
return;
|
||||
}
|
||||
res.status(502).json({
|
||||
message: sanitizeUserFacingProxyMessage(
|
||||
err instanceof Error ? err.message : 'SSE 代理失败',
|
||||
@@ -1467,6 +1706,7 @@ export function createTkmindProxy({
|
||||
proxySessionEvents,
|
||||
resolveTarget,
|
||||
startSessionForUser,
|
||||
getRuntimeStatus,
|
||||
submitSessionReplyForUser,
|
||||
apiFetch: async (pathname, init) => apiFetch(await pickTarget(), apiSecret, pathname, init),
|
||||
apiFetchTo: (target, pathname, init) => apiFetch(target, apiSecret, pathname, init),
|
||||
|
||||
Reference in New Issue
Block a user