Files
memind/agent-run-gateway.mjs
T
2026-07-02 12:38:34 +08:00

733 lines
24 KiB
JavaScript

import crypto from 'node:crypto';
import fs from 'node:fs/promises';
import path from 'node:path';
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;
const DEFAULT_RUN_HEARTBEAT_MS = 30 * 1000;
const TOOL_GATEWAY_SUMMARY_LIMIT = 4096;
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 summarizeText(value, limit = TOOL_GATEWAY_SUMMARY_LIMIT) {
const text = String(value ?? '');
if (text.length <= limit) return text;
return text.slice(text.length - limit);
}
function normalizeExpectedFileCheck(value) {
if (typeof value === 'string') {
const expectedPath = value.trim();
return expectedPath ? { path: expectedPath } : null;
}
if (!value || typeof value !== 'object' || Array.isArray(value)) return null;
const expectedPath = String(value.path ?? value.file ?? value.relativePath ?? '').trim();
if (!expectedPath) return null;
const contains = value.contains ?? value.expectedContent ?? value.contentIncludes;
return {
path: expectedPath,
...(contains == null ? {} : { contains: String(contains) }),
};
}
function normalizeToolGatewayValidation(value) {
if (!value || typeof value !== 'object' || Array.isArray(value)) return null;
const checks = [];
const single = normalizeExpectedFileCheck(value.expectedFile ?? value.expectedPath);
if (single) checks.push(single);
if (Array.isArray(value.expectedFiles)) {
for (const item of value.expectedFiles) {
const check = normalizeExpectedFileCheck(item);
if (check) checks.push(check);
}
}
return checks.length > 0 ? { expectedFiles: checks } : null;
}
function resolveValidationPath(cwd, expectedPath) {
if (!String(cwd ?? '').trim()) {
const err = new Error('Tool Gateway validation failed: missing working directory');
err.code = 'TOOL_GATEWAY_VALIDATION_FAILED';
err.retryable = false;
throw err;
}
const root = path.resolve(String(cwd ?? ''));
const target = path.resolve(root, String(expectedPath ?? ''));
if (target !== root && !target.startsWith(`${root}${path.sep}`)) {
const err = new Error(`Tool Gateway validation path escapes working directory: ${expectedPath}`);
err.code = 'TOOL_GATEWAY_VALIDATION_FAILED';
err.retryable = false;
throw err;
}
return { root, target };
}
async function validateToolGatewayResult({ result, validation, cwd }) {
if (!validation?.expectedFiles?.length || result?.dryRun) {
return null;
}
const checks = [];
for (const expected of validation.expectedFiles) {
const { root, target } = resolveValidationPath(cwd ?? result?.cwd, expected.path);
let content = '';
let stat = null;
try {
stat = await fs.stat(target);
content = await fs.readFile(target, 'utf8');
} catch (cause) {
const err = new Error(`Tool Gateway validation failed: expected file not found: ${expected.path}`);
err.code = 'TOOL_GATEWAY_VALIDATION_FAILED';
err.retryable = false;
err.validation = {
expectedFile: expected.path,
cwd: root,
reason: 'missing_file',
};
err.cause = cause;
throw err;
}
if (expected.contains != null && !content.includes(expected.contains)) {
const err = new Error(`Tool Gateway validation failed: expected content not found in ${expected.path}`);
err.code = 'TOOL_GATEWAY_VALIDATION_FAILED';
err.retryable = false;
err.validation = {
expectedFile: expected.path,
cwd: root,
reason: 'missing_content',
};
throw err;
}
checks.push({
path: expected.path,
sizeBytes: Number(stat?.size ?? 0),
contains: expected.contains == null ? null : true,
});
}
return { expectedFiles: checks };
}
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),
validation: normalizeToolGatewayValidation(runMetadata?.validation ?? metadata?.toolGatewayValidation),
};
}
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,
toolGateway = null,
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,
),
heartbeatMs = positiveInteger(
process.env.MEMIND_AGENT_RUN_HEARTBEAT_MS,
DEFAULT_RUN_HEARTBEAT_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);
}
function startRunHeartbeat(runId, { attempt }) {
let stopped = false;
const writeHeartbeat = async () => {
if (stopped) return;
try {
await appendEvent(runId, 'worker_heartbeat', {
attempt,
pid: process.pid,
heartbeatMs,
});
} catch (err) {
console.error('[AgentRun] worker heartbeat failed:', err instanceof Error ? err.message : err);
}
};
void writeHeartbeat();
const timer = setInterval(() => {
void writeHeartbeat();
}, heartbeatMs);
timer.unref?.();
return () => {
stopped = true;
clearInterval(timer);
};
}
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);
const toolGatewayStatus = toolGateway?.getStatus ? toolGateway.getStatus() : null;
if (runOptions.toolMode === 'code' && toolGatewayStatus?.enabled) {
const workingDir = userAuth?.resolveWorkingDir
? await userAuth.resolveWorkingDir(row.user_id)
: undefined;
await appendEvent(runId, 'tool_gateway_dispatch', {
protocol: toolGatewayStatus.protocol ?? 'agent-run-v1',
taskType: runOptions.taskType,
workingDir: workingDir ?? null,
});
const result = await toolGateway.executeJob({
runId,
userId: row.user_id,
requestId: row.request_id,
userMessage,
taskType: runOptions.taskType,
cwd: workingDir,
timeoutMs: runTimeoutMs,
});
await appendEvent(runId, 'tool_gateway_result', {
executor: result.executor ?? null,
dryRun: Boolean(result.dryRun),
exitCode: result.exitCode ?? null,
stdoutTail: summarizeText(result.stdout),
stderrTail: summarizeText(result.stderr),
});
try {
const validation = await validateToolGatewayResult({
result,
validation: runOptions.validation,
cwd: workingDir,
});
if (validation) {
await appendEvent(runId, 'tool_gateway_validation', validation);
}
} catch (err) {
await appendEvent(runId, 'tool_gateway_validation_failed', {
code: err?.code ?? null,
message: err instanceof Error ? err.message : String(err),
validation: err?.validation ?? null,
});
throw err;
}
return { sessionId: row.agent_session_id ?? null };
}
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 });
const stopHeartbeat = startRunHeartbeat(runId, { 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 && err?.retryable !== false && 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]);
}
} finally {
stopHeartbeat();
}
}
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);
}
const [oldestRunningRows] = await pool.query(
`SELECT
r.id,
r.request_id,
r.started_at,
r.updated_at,
r.attempts,
h.latest_heartbeat_at
FROM h5_agent_runs r
LEFT JOIN (
SELECT run_id, MAX(created_at) AS latest_heartbeat_at
FROM h5_agent_run_events
WHERE event_type = 'worker_heartbeat'
GROUP BY run_id
) h ON h.run_id = r.id
WHERE r.status = 'running'
ORDER BY COALESCE(h.latest_heartbeat_at, r.started_at) ASC
LIMIT 1`,
);
const oldestRunningStartedAt = oldestRunningRows[0]?.started_at == null
? null
: Number(oldestRunningRows[0].started_at);
const oldestRunningAgeMs = oldestRunningStartedAt == null
? 0
: Math.max(0, nowMs() - oldestRunningStartedAt);
const oldestRunningHeartbeatAt = oldestRunningRows[0]?.latest_heartbeat_at == null
? null
: Number(oldestRunningRows[0].latest_heartbeat_at);
const oldestRunningHeartbeatAgeMs = oldestRunningRows[0]
? Math.max(0, nowMs() - Number(oldestRunningRows[0].latest_heartbeat_at ?? oldestRunningRows[0].started_at ?? nowMs()))
: 0;
const [runningWithoutHeartbeatRows] = await pool.query(
`SELECT COUNT(*) AS count
FROM h5_agent_runs r
LEFT JOIN (
SELECT run_id, MAX(created_at) AS latest_heartbeat_at
FROM h5_agent_run_events
WHERE event_type = 'worker_heartbeat'
GROUP BY run_id
) h ON h.run_id = r.id
WHERE r.status = 'running' AND h.latest_heartbeat_at IS NULL`,
);
return {
autoDispatch,
maxConcurrentRuns,
runTimeoutMs,
heartbeatMs,
inFlight: inFlight.size,
pendingDispatches: queuedDispatches.length,
statusCounts,
oldestRunningStartedAt,
oldestRunningAgeMs,
oldestRunningHeartbeatAt,
oldestRunningHeartbeatAgeMs,
runningWithoutHeartbeatCount: Number(runningWithoutHeartbeatRows[0]?.count ?? 0),
latestRunningRun: oldestRunningRows[0] ? {
id: oldestRunningRows[0].id,
requestId: oldestRunningRows[0].request_id,
startedAt: oldestRunningStartedAt,
updatedAt: oldestRunningRows[0].updated_at == null ? null : Number(oldestRunningRows[0].updated_at),
attempts: Number(oldestRunningRows[0].attempts ?? 0),
heartbeatAt: oldestRunningHeartbeatAt,
heartbeatAgeMs: oldestRunningHeartbeatAgeMs,
} : null,
terminalStatuses: [...TERMINAL_STATUSES],
toolGateway: toolGateway?.getStatus ? toolGateway.getStatus() : null,
};
}
async function recoverStaleRunningRuns({
staleMs = runTimeoutMs,
limit = maxConcurrentRuns,
dryRun = true,
reason = 'stale_running_timeout',
} = {}) {
const normalizedStaleMs = positiveInteger(staleMs, runTimeoutMs);
const normalizedLimit = positiveInteger(limit, maxConcurrentRuns);
const cutoff = nowMs() - normalizedStaleMs;
const [rows] = await pool.query(
`SELECT
r.id,
r.request_id,
r.started_at,
r.updated_at,
r.attempts,
h.latest_heartbeat_at
FROM h5_agent_runs r
LEFT JOIN (
SELECT run_id, MAX(created_at) AS latest_heartbeat_at
FROM h5_agent_run_events
WHERE event_type = 'worker_heartbeat'
GROUP BY run_id
) h ON h.run_id = r.id
WHERE r.status = 'running'
AND r.started_at IS NOT NULL
AND COALESCE(h.latest_heartbeat_at, r.started_at) <= ?
ORDER BY COALESCE(h.latest_heartbeat_at, r.started_at) ASC
LIMIT ?`,
[cutoff, normalizedLimit],
);
const recovered = [];
for (const row of rows) {
const ageMs = Math.max(0, nowMs() - Number(row.started_at ?? 0));
const heartbeatAt = row.latest_heartbeat_at == null ? null : Number(row.latest_heartbeat_at);
const heartbeatAgeMs = Math.max(0, nowMs() - Number(row.latest_heartbeat_at ?? row.started_at ?? 0));
const item = {
id: row.id,
requestId: row.request_id,
startedAt: Number(row.started_at ?? 0),
updatedAt: Number(row.updated_at ?? 0),
attempts: Number(row.attempts ?? 0),
heartbeatAt,
heartbeatAgeMs,
ageMs,
reason,
};
if (!dryRun) {
const message = heartbeatAt == null
? `agent run recovered from stale running state after ${ageMs}ms without heartbeat`
: `agent run recovered from stale running state after heartbeat was stale for ${heartbeatAgeMs}ms`;
const completedAt = nowMs();
const [update] = await pool.query(
`UPDATE h5_agent_runs
SET status = 'failed', error_message = ?, updated_at = ?, completed_at = ?
WHERE id = ?
AND status = 'running'
AND started_at IS NOT NULL
AND COALESCE(
(SELECT MAX(created_at)
FROM h5_agent_run_events
WHERE run_id = h5_agent_runs.id AND event_type = 'worker_heartbeat'),
started_at
) <= ?`,
[message, completedAt, completedAt, row.id, cutoff],
);
if (Number(update?.affectedRows ?? 0) === 0) continue;
await appendEvent(row.id, 'stale_recovered', {
reason,
staleMs: normalizedStaleMs,
ageMs,
heartbeatAt,
heartbeatAgeMs,
status: 'failed',
error: message,
});
}
recovered.push(item);
}
return {
dryRun,
staleMs: normalizedStaleMs,
cutoff,
considered: rows.length,
recovered: dryRun ? 0 : recovered.length,
runs: recovered,
};
}
async function dispatchQueuedRuns({ limit = maxConcurrentRuns } = {}) {
const staleRecovery = await recoverStaleRunningRuns({
staleMs: runTimeoutMs,
limit: Math.max(maxConcurrentRuns, Number(limit) || maxConcurrentRuns),
dryRun: false,
reason: 'worker_dispatch_stale_recovery',
});
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,
staleRecovery,
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,
staleRecovery,
queue: await getQueueStatus(),
};
}
return {
createRun,
getRunForUser,
dispatchRun,
getQueueStatus,
dispatchQueuedRuns,
recoverStaleRunningRuns,
};
}