John
396 lines
12 KiB
JavaScript
396 lines
12 KiB
JavaScript
import crypto from 'node:crypto';
|
|
|
|
const DEFAULT_RUN_RETRY_DELAYS_MS = [1500, 5000, 15000];
|
|
const TERMINAL_STATUSES = new Set(['succeeded', 'failed']);
|
|
const CODE_TOOL_MODES = new Set(['code', 'code-task', 'code_task', 'code-tool', 'code_tool', 'code_tool_task']);
|
|
const RUN_METADATA_KEY = 'memindRun';
|
|
const DEFAULT_MAX_CONCURRENT_RUNS = 1;
|
|
const DEFAULT_RUN_TIMEOUT_MS = 15 * 60 * 1000;
|
|
|
|
function nowMs() {
|
|
return Date.now();
|
|
}
|
|
|
|
function safeJsonParse(value, fallback = null) {
|
|
try {
|
|
return JSON.parse(value);
|
|
} catch {
|
|
return fallback;
|
|
}
|
|
}
|
|
|
|
function serializeMessage(message) {
|
|
return JSON.stringify(message ?? {});
|
|
}
|
|
|
|
function positiveInteger(value, fallback) {
|
|
const n = Number(value);
|
|
if (!Number.isFinite(n) || n <= 0) return fallback;
|
|
return Math.floor(n);
|
|
}
|
|
|
|
function timeoutError(ms) {
|
|
const err = new Error(`agent run timed out after ${ms}ms`);
|
|
err.code = 'AGENT_RUN_TIMEOUT';
|
|
return err;
|
|
}
|
|
|
|
function envFlag(value, fallback = false) {
|
|
const raw = String(value ?? '').trim().toLowerCase();
|
|
if (!raw) return fallback;
|
|
return ['1', 'true', 'yes', 'on'].includes(raw);
|
|
}
|
|
|
|
export function normalizeAgentRunToolMode(value) {
|
|
const normalized = String(value ?? 'chat').trim().toLowerCase();
|
|
if (!normalized || normalized === 'chat') return 'chat';
|
|
if (CODE_TOOL_MODES.has(normalized)) return 'code';
|
|
throw new Error(`不支持的 tool_mode: ${value}`);
|
|
}
|
|
|
|
function normalizeTaskType(value) {
|
|
const normalized = String(value ?? '').trim();
|
|
return normalized || null;
|
|
}
|
|
|
|
function withRunMetadata(userMessage, { toolMode = 'chat', taskType = null } = {}) {
|
|
const message = (userMessage && typeof userMessage === 'object' && !Array.isArray(userMessage))
|
|
? { ...userMessage }
|
|
: { value: userMessage };
|
|
const metadata = (message.metadata && typeof message.metadata === 'object' && !Array.isArray(message.metadata))
|
|
? { ...message.metadata }
|
|
: {};
|
|
metadata[RUN_METADATA_KEY] = {
|
|
...(metadata[RUN_METADATA_KEY] && typeof metadata[RUN_METADATA_KEY] === 'object'
|
|
? metadata[RUN_METADATA_KEY]
|
|
: {}),
|
|
toolMode: normalizeAgentRunToolMode(toolMode),
|
|
...(taskType ? { taskType } : {}),
|
|
};
|
|
return { ...message, metadata };
|
|
}
|
|
|
|
function getRunOptionsFromMessage(userMessage) {
|
|
const metadata = userMessage?.metadata;
|
|
const runMetadata = metadata?.[RUN_METADATA_KEY] ?? metadata?.agentRun ?? {};
|
|
let toolMode = 'chat';
|
|
try {
|
|
toolMode = normalizeAgentRunToolMode(runMetadata?.toolMode ?? metadata?.toolMode ?? 'chat');
|
|
} catch {
|
|
toolMode = 'chat';
|
|
}
|
|
return {
|
|
toolMode,
|
|
taskType: normalizeTaskType(runMetadata?.taskType ?? metadata?.taskType),
|
|
};
|
|
}
|
|
|
|
function projectRun(row) {
|
|
if (!row) return null;
|
|
return {
|
|
id: row.id,
|
|
userId: row.user_id,
|
|
sessionId: row.agent_session_id ?? null,
|
|
requestId: row.request_id,
|
|
status: row.status,
|
|
attempts: Number(row.attempts ?? 0),
|
|
error: row.error_message ?? null,
|
|
createdAt: Number(row.created_at ?? 0),
|
|
updatedAt: Number(row.updated_at ?? 0),
|
|
startedAt: row.started_at == null ? null : Number(row.started_at),
|
|
completedAt: row.completed_at == null ? null : Number(row.completed_at),
|
|
};
|
|
}
|
|
|
|
export function createAgentRunGateway({
|
|
pool,
|
|
userAuth,
|
|
tkmindProxy,
|
|
retryDelaysMs = DEFAULT_RUN_RETRY_DELAYS_MS,
|
|
autoDispatch = envFlag(process.env.MEMIND_AGENT_RUN_AUTODISPATCH, true),
|
|
maxConcurrentRuns = positiveInteger(
|
|
process.env.MEMIND_AGENT_RUN_QUEUE_CONCURRENCY,
|
|
DEFAULT_MAX_CONCURRENT_RUNS,
|
|
),
|
|
runTimeoutMs = positiveInteger(
|
|
process.env.MEMIND_AGENT_RUN_TIMEOUT_MS,
|
|
DEFAULT_RUN_TIMEOUT_MS,
|
|
),
|
|
}) {
|
|
const inFlight = new Set();
|
|
const queuedDispatches = [];
|
|
const queuedDispatchSet = new Set();
|
|
|
|
function enqueueRun(runId) {
|
|
if (!runId || inFlight.has(runId) || queuedDispatchSet.has(runId)) return false;
|
|
queuedDispatchSet.add(runId);
|
|
queuedDispatches.push(runId);
|
|
return true;
|
|
}
|
|
|
|
function nextQueuedRunId() {
|
|
const runId = queuedDispatches.shift();
|
|
if (runId) queuedDispatchSet.delete(runId);
|
|
return runId ?? null;
|
|
}
|
|
|
|
async function appendEvent(runId, type, data = null) {
|
|
await pool.query(
|
|
`INSERT INTO h5_agent_run_events (id, run_id, event_type, data_json, created_at)
|
|
VALUES (?, ?, ?, ?, ?)`,
|
|
[
|
|
crypto.randomUUID(),
|
|
runId,
|
|
type,
|
|
data == null ? null : JSON.stringify(data),
|
|
nowMs(),
|
|
],
|
|
);
|
|
}
|
|
|
|
async function getRunById(runId) {
|
|
const [rows] = await pool.query(
|
|
`SELECT * FROM h5_agent_runs WHERE id = ? LIMIT 1`,
|
|
[runId],
|
|
);
|
|
return rows[0] ?? null;
|
|
}
|
|
|
|
async function getRunForUser(userId, runId) {
|
|
const [rows] = await pool.query(
|
|
`SELECT * FROM h5_agent_runs WHERE id = ? AND user_id = ? LIMIT 1`,
|
|
[runId, userId],
|
|
);
|
|
return projectRun(rows[0] ?? null);
|
|
}
|
|
|
|
async function getRunByRequest(userId, requestId) {
|
|
const [rows] = await pool.query(
|
|
`SELECT * FROM h5_agent_runs WHERE user_id = ? AND request_id = ? LIMIT 1`,
|
|
[userId, requestId],
|
|
);
|
|
return rows[0] ?? null;
|
|
}
|
|
|
|
async function createRun(userId, {
|
|
sessionId = null,
|
|
requestId,
|
|
userMessage,
|
|
toolMode = 'chat',
|
|
taskType = null,
|
|
}) {
|
|
const normalizedRequestId = String(requestId ?? '').trim();
|
|
if (!normalizedRequestId) {
|
|
throw new Error('缺少 request_id');
|
|
}
|
|
const normalizedToolMode = normalizeAgentRunToolMode(toolMode);
|
|
const normalizedTaskType = normalizeTaskType(taskType);
|
|
const existing = await getRunByRequest(userId, normalizedRequestId);
|
|
if (existing) {
|
|
if (autoDispatch && !TERMINAL_STATUSES.has(existing.status)) dispatchRun(existing.id);
|
|
return projectRun(existing);
|
|
}
|
|
|
|
const runId = crypto.randomUUID();
|
|
const createdAt = nowMs();
|
|
const runMessage = withRunMetadata(userMessage, {
|
|
toolMode: normalizedToolMode,
|
|
taskType: normalizedTaskType,
|
|
});
|
|
await pool.query(
|
|
`INSERT INTO h5_agent_runs
|
|
(id, user_id, agent_session_id, request_id, status, attempts,
|
|
user_message_json, error_message, created_at, updated_at, started_at, completed_at)
|
|
VALUES (?, ?, ?, ?, 'queued', 0, ?, NULL, ?, ?, NULL, NULL)`,
|
|
[
|
|
runId,
|
|
userId,
|
|
sessionId || null,
|
|
normalizedRequestId,
|
|
serializeMessage(runMessage),
|
|
createdAt,
|
|
createdAt,
|
|
],
|
|
);
|
|
await appendEvent(runId, 'queued', {
|
|
sessionId: sessionId || null,
|
|
toolMode: normalizedToolMode,
|
|
taskType: normalizedTaskType,
|
|
});
|
|
if (autoDispatch) dispatchRun(runId);
|
|
return projectRun(await getRunById(runId));
|
|
}
|
|
|
|
async function markRun(runId, status, fields = {}) {
|
|
const updates = ['status = ?', 'updated_at = ?'];
|
|
const values = [status, nowMs()];
|
|
for (const [key, value] of Object.entries(fields)) {
|
|
updates.push(`${key} = ?`);
|
|
values.push(value);
|
|
}
|
|
values.push(runId);
|
|
await pool.query(
|
|
`UPDATE h5_agent_runs SET ${updates.join(', ')} WHERE id = ?`,
|
|
values,
|
|
);
|
|
await appendEvent(runId, status, fields);
|
|
}
|
|
|
|
async function runWithTimeout(runId, task) {
|
|
let timer = null;
|
|
const timeout = new Promise((_, reject) => {
|
|
timer = setTimeout(() => reject(timeoutError(runTimeoutMs)), runTimeoutMs);
|
|
});
|
|
try {
|
|
return await Promise.race([task(), timeout]);
|
|
} finally {
|
|
if (timer) clearTimeout(timer);
|
|
}
|
|
}
|
|
|
|
async function executeRun(row, runId) {
|
|
const userMessage = safeJsonParse(row.user_message_json, {});
|
|
const runOptions = getRunOptionsFromMessage(userMessage);
|
|
let sessionId = row.agent_session_id ?? null;
|
|
if (!sessionId) {
|
|
const sessionOptions = {};
|
|
if (runOptions.toolMode === 'code' && userAuth?.getCodeAgentSessionPolicy) {
|
|
sessionOptions.sessionPolicy = await userAuth.getCodeAgentSessionPolicy(row.user_id);
|
|
}
|
|
const session = await tkmindProxy.startSessionForUser(row.user_id, sessionOptions);
|
|
sessionId = session.id;
|
|
await pool.query(
|
|
`UPDATE h5_agent_runs SET agent_session_id = ?, updated_at = ? WHERE id = ?`,
|
|
[sessionId, nowMs(), runId],
|
|
);
|
|
await appendEvent(runId, 'session_started', {
|
|
sessionId,
|
|
toolMode: runOptions.toolMode,
|
|
taskType: runOptions.taskType,
|
|
});
|
|
}
|
|
|
|
await tkmindProxy.submitSessionReplyForUser(
|
|
row.user_id,
|
|
sessionId,
|
|
row.request_id,
|
|
userMessage,
|
|
{ toolMode: runOptions.toolMode },
|
|
);
|
|
return { sessionId };
|
|
}
|
|
|
|
async function processRun(runId) {
|
|
const row = await getRunById(runId);
|
|
if (!row || TERMINAL_STATUSES.has(row.status)) return;
|
|
const nextAttempt = Number(row.attempts ?? 0) + 1;
|
|
const [claim] = await pool.query(
|
|
`UPDATE h5_agent_runs
|
|
SET status = 'running', attempts = ?, started_at = COALESCE(started_at, ?), updated_at = ?, error_message = NULL
|
|
WHERE id = ? AND status IN ('queued', 'retryable')`,
|
|
[nextAttempt, nowMs(), nowMs(), runId],
|
|
);
|
|
if (Number(claim?.affectedRows ?? 0) === 0) return;
|
|
await appendEvent(runId, 'running', { attempt: nextAttempt });
|
|
|
|
try {
|
|
const { sessionId } = await runWithTimeout(runId, () => executeRun(row, runId));
|
|
await markRun(runId, 'succeeded', {
|
|
agent_session_id: sessionId,
|
|
completed_at: nowMs(),
|
|
});
|
|
} catch (err) {
|
|
const message = err instanceof Error ? err.message : String(err);
|
|
const timedOut = err?.code === 'AGENT_RUN_TIMEOUT';
|
|
if (timedOut) {
|
|
await appendEvent(runId, 'timeout', { timeoutMs: runTimeoutMs });
|
|
}
|
|
const retryable = !timedOut && nextAttempt < retryDelaysMs.length;
|
|
await markRun(runId, retryable ? 'retryable' : 'failed', {
|
|
error_message: message,
|
|
completed_at: retryable ? null : nowMs(),
|
|
});
|
|
if (retryable && autoDispatch) {
|
|
setTimeout(() => dispatchRun(runId), retryDelaysMs[nextAttempt - 1]);
|
|
}
|
|
}
|
|
}
|
|
|
|
function drainQueue() {
|
|
while (inFlight.size < maxConcurrentRuns) {
|
|
const runId = nextQueuedRunId();
|
|
if (!runId) break;
|
|
inFlight.add(runId);
|
|
void processRun(runId)
|
|
.finally(() => {
|
|
inFlight.delete(runId);
|
|
drainQueue();
|
|
});
|
|
}
|
|
}
|
|
|
|
function dispatchRun(runId) {
|
|
if (!enqueueRun(runId)) return;
|
|
drainQueue();
|
|
}
|
|
|
|
async function getQueueStatus() {
|
|
const [rows] = await pool.query(
|
|
`SELECT status, COUNT(*) AS count
|
|
FROM h5_agent_runs
|
|
WHERE status IN ('queued', 'running', 'retryable')
|
|
GROUP BY status`,
|
|
);
|
|
const statusCounts = {};
|
|
for (const row of rows) {
|
|
statusCounts[row.status] = Number(row.count ?? 0);
|
|
}
|
|
return {
|
|
autoDispatch,
|
|
maxConcurrentRuns,
|
|
runTimeoutMs,
|
|
inFlight: inFlight.size,
|
|
pendingDispatches: queuedDispatches.length,
|
|
statusCounts,
|
|
terminalStatuses: [...TERMINAL_STATUSES],
|
|
};
|
|
}
|
|
|
|
async function dispatchQueuedRuns({ limit = maxConcurrentRuns } = {}) {
|
|
const dispatchLimit = Math.max(1, Number(limit) || maxConcurrentRuns);
|
|
const available = Math.max(0, maxConcurrentRuns - inFlight.size - queuedDispatches.length);
|
|
if (available <= 0) {
|
|
return {
|
|
considered: 0,
|
|
dispatched: 0,
|
|
queue: await getQueueStatus(),
|
|
};
|
|
}
|
|
const take = Math.min(dispatchLimit, available);
|
|
const [rows] = await pool.query(
|
|
`SELECT id
|
|
FROM h5_agent_runs
|
|
WHERE status IN ('queued', 'retryable')
|
|
ORDER BY updated_at ASC
|
|
LIMIT ?`,
|
|
[take],
|
|
);
|
|
for (const row of rows) {
|
|
dispatchRun(row.id);
|
|
}
|
|
return {
|
|
considered: rows.length,
|
|
dispatched: rows.length,
|
|
queue: await getQueueStatus(),
|
|
};
|
|
}
|
|
|
|
return {
|
|
createRun,
|
|
getRunForUser,
|
|
dispatchRun,
|
|
getQueueStatus,
|
|
dispatchQueuedRuns,
|
|
};
|
|
}
|