feat: add agent run queue controls

This commit is contained in:
John
2026-07-02 07:52:07 +08:00
parent 8e5094fcd9
commit 3ea973968e
9 changed files with 387 additions and 66 deletions
+4
View File
@@ -26,6 +26,10 @@ H5_PUBLIC_BASE_URL=http://127.0.0.1:5173
# VITE_AGENT_CODE_RUNS_ENABLED=0
# VITE_AGENT_CODE_RUNS_AUTODETECT=0
# Tool Gateway Queue v0/agent/runs 后台任务队列,默认保守限流)。
# MEMIND_AGENT_RUN_QUEUE_CONCURRENCY=1
# MEMIND_AGENT_RUN_TIMEOUT_MS=900000
# ===== memind_adm(独立后台服务)=====
# 平台超管 /admin-api/* 与广场运营 /api/ops/v1/* 已从公网服务拆出,独立进程运行。
# 登录仍在主域完成;后台凭 H5_COOKIE_DOMAIN=.tkmind.cn 共享的会话 Cookie 鉴权(生产务必设置)。
+1
View File
@@ -46,3 +46,4 @@ Streaming runtime operations:
node scripts/runtime-worker-drain.mjs undrain goosed-3
node scripts/check-tool-runtime.mjs
node scripts/check-agent-code-run-entry.mjs
curl -sk https://mm.tkmind.cn/api/runtime/status # includes toolRuntime.queue
@@ -148,6 +148,7 @@ function summarizeRuntime(runtimeJson) {
routerEnabled: Boolean(runtimeJson?.router?.enabled),
publicBaseUrl: runtimeJson?.publicBaseUrl ?? null,
toolRuntime: runtimeJson?.toolRuntime ?? null,
toolQueue: runtimeJson?.toolRuntime?.queue ?? null,
workers: workers.map((worker) => ({
id: worker.id,
healthy: runtimeJson?.targets?.find((target) => target.target === worker.target)?.healthy ?? null,
@@ -220,6 +221,10 @@ for (const worker of runtimeSummary?.workers ?? []) {
if (worker.firstTokenCount > 0 && worker.ewmaFirstTokenMs <= 0) failures.push(`${worker.id}_first_token_ewma_missing`);
}
if (runtimeSummary?.toolRuntime?.chatInjectsCodeTools !== false) failures.push('chat_injects_code_tools');
if (runtimeSummary?.toolQueue?.error) failures.push('tool_queue_status_error');
if (runtimeSummary?.toolQueue?.inFlight > runtimeSummary?.toolQueue?.maxConcurrentRuns) {
failures.push('tool_queue_concurrency_exceeded');
}
const report = {
ok: failures.length === 0,
+128 -34
View File
@@ -6509,6 +6509,8 @@ var DEFAULT_RUN_RETRY_DELAYS_MS = [1500, 5e3, 15e3];
var TERMINAL_STATUSES = /* @__PURE__ */ new Set(["succeeded", "failed"]);
var CODE_TOOL_MODES = /* @__PURE__ */ new Set(["code", "code-task", "code_task", "code-tool", "code_tool", "code_tool_task"]);
var RUN_METADATA_KEY = "memindRun";
var DEFAULT_MAX_CONCURRENT_RUNS = 1;
var DEFAULT_RUN_TIMEOUT_MS = 15 * 60 * 1e3;
function nowMs() {
return Date.now();
}
@@ -6522,6 +6524,16 @@ function safeJsonParse(value, fallback = null) {
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 normalizeAgentRunToolMode(value) {
const normalized = String(value ?? "chat").trim().toLowerCase();
if (!normalized || normalized === "chat") return "chat";
@@ -6577,9 +6589,30 @@ function createAgentRunGateway({
userAuth: userAuth2,
tkmindProxy: tkmindProxy2,
retryDelaysMs = DEFAULT_RUN_RETRY_DELAYS_MS,
autoDispatch = true
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 = /* @__PURE__ */ new Set();
const queuedDispatches = [];
const queuedDispatchSet = /* @__PURE__ */ 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 pool2.query(
`INSERT INTO h5_agent_run_events (id, run_id, event_type, data_json, created_at)
@@ -6675,6 +6708,47 @@ function createAgentRunGateway({
);
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" && userAuth2?.getCodeAgentSessionPolicy) {
sessionOptions.sessionPolicy = await userAuth2.getCodeAgentSessionPolicy(row.user_id);
}
const session = await tkmindProxy2.startSessionForUser(row.user_id, sessionOptions);
sessionId = session.id;
await pool2.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 tkmindProxy2.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;
@@ -6688,40 +6762,18 @@ function createAgentRunGateway({
if (Number(claim?.affectedRows ?? 0) === 0) return;
await appendEvent(runId, "running", { attempt: nextAttempt });
try {
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" && userAuth2?.getCodeAgentSessionPolicy) {
sessionOptions.sessionPolicy = await userAuth2.getCodeAgentSessionPolicy(row.user_id);
}
const session = await tkmindProxy2.startSessionForUser(row.user_id, sessionOptions);
sessionId = session.id;
await pool2.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 tkmindProxy2.submitSessionReplyForUser(
row.user_id,
sessionId,
row.request_id,
userMessage,
{ toolMode: runOptions.toolMode }
);
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 retryable = nextAttempt < retryDelaysMs.length;
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()
@@ -6731,15 +6783,46 @@ function createAgentRunGateway({
}
}
}
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 (!runId || inFlight.has(runId)) return;
inFlight.add(runId);
void processRun(runId).finally(() => inFlight.delete(runId));
if (!enqueueRun(runId)) return;
drainQueue();
}
async function getQueueStatus() {
const [rows] = await pool2.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 {
maxConcurrentRuns,
runTimeoutMs,
inFlight: inFlight.size,
pendingDispatches: queuedDispatches.length,
statusCounts,
terminalStatuses: [...TERMINAL_STATUSES]
};
}
return {
createRun,
getRunForUser,
dispatchRun
dispatchRun,
getQueueStatus
};
}
@@ -34957,7 +35040,9 @@ async function bootstrapUserAuth() {
agentRunGateway = createAgentRunGateway({
pool: pool2,
userAuth,
tkmindProxy
tkmindProxy,
maxConcurrentRuns: Number(process.env.MEMIND_AGENT_RUN_QUEUE_CONCURRENCY ?? 1),
runTimeoutMs: Number(process.env.MEMIND_AGENT_RUN_TIMEOUT_MS ?? 15 * 60 * 1e3)
});
wechatMpService = createWechatMpService({
config: WECHAT_MP_CONFIG,
@@ -36043,6 +36128,15 @@ api.get("/runtime/status", async (_req, res) => {
}
try {
const status = await tkmindProxy.getRuntimeStatus();
const toolQueue = agentRunGateway?.getQueueStatus ? await agentRunGateway.getQueueStatus().catch((err) => ({
error: err instanceof Error ? err.message : String(err)
})) : null;
if (toolQueue) {
status.toolRuntime = {
...status.toolRuntime ?? {},
queue: toolQueue
};
}
return res.json({
ok: true,
timestamp: (/* @__PURE__ */ new Date()).toISOString(),
+124 -32
View File
@@ -4,6 +4,8 @@ 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();
@@ -21,6 +23,18 @@ 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;
}
export function normalizeAgentRunToolMode(value) {
const normalized = String(value ?? 'chat').trim().toLowerCase();
if (!normalized || normalized === 'chat') return 'chat';
@@ -88,8 +102,31 @@ export function createAgentRunGateway({
tkmindProxy,
retryDelaysMs = DEFAULT_RUN_RETRY_DELAYS_MS,
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(
@@ -193,6 +230,50 @@ export function createAgentRunGateway({
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;
@@ -207,41 +288,18 @@ export function createAgentRunGateway({
await appendEvent(runId, 'running', { attempt: nextAttempt });
try {
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 },
);
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 retryable = nextAttempt < retryDelaysMs.length;
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(),
@@ -252,15 +310,49 @@ export function createAgentRunGateway({
}
}
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 (!runId || inFlight.has(runId)) return;
inFlight.add(runId);
void processRun(runId).finally(() => inFlight.delete(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 {
maxConcurrentRuns,
runTimeoutMs,
inFlight: inFlight.size,
pendingDispatches: queuedDispatches.length,
statusCounts,
terminalStatuses: [...TERMINAL_STATUSES],
};
}
return {
createRun,
getRunForUser,
dispatchRun,
getQueueStatus,
};
}
+106
View File
@@ -22,6 +22,14 @@ function createFakePool() {
if (sql.includes('SELECT * FROM h5_agent_runs WHERE id = ? LIMIT 1')) {
return [[runs.get(params[0])].filter(Boolean)];
}
if (sql.includes('SELECT status, COUNT(*) AS count')) {
const counts = new Map();
for (const row of runs.values()) {
if (!['queued', 'running', 'retryable'].includes(row.status)) continue;
counts.set(row.status, (counts.get(row.status) ?? 0) + 1);
}
return [[...counts].map(([status, count]) => ({ status, count }))];
}
if (sql.includes('INSERT INTO h5_agent_runs')) {
const [
id,
@@ -232,3 +240,101 @@ test('agent run retries transient failures and then becomes terminal', async ()
assert.equal(pool.runs.get(run.id).attempts, 2);
assert.match(pool.runs.get(run.id).error_message, /upstream unavailable/);
});
test('agent run queue limits concurrent execution', async () => {
const pool = createFakePool();
let active = 0;
let maxActive = 0;
const release = [];
const gateway = createAgentRunGateway({
pool,
userAuth: {},
tkmindProxy: {
async startSessionForUser(_userId) {
return { id: `session-${release.length + 1}` };
},
async submitSessionReplyForUser() {
active += 1;
maxActive = Math.max(maxActive, active);
await new Promise((resolve) => release.push(resolve));
active -= 1;
},
},
retryDelaysMs: [],
maxConcurrentRuns: 1,
});
const run1 = await gateway.createRun('user-1', {
requestId: 'req-1',
userMessage: { role: 'user', content: [] },
});
const run2 = await gateway.createRun('user-1', {
requestId: 'req-2',
userMessage: { role: 'user', content: [] },
});
await waitFor(() => active === 1 && pool.runs.get(run2.id)?.status === 'queued');
assert.equal(maxActive, 1);
assert.equal((await gateway.getQueueStatus()).pendingDispatches, 1);
release.shift()();
await waitFor(() => pool.runs.get(run1.id)?.status === 'succeeded' && active === 1);
release.shift()();
await waitFor(() => pool.runs.get(run2.id)?.status === 'succeeded');
assert.equal(maxActive, 1);
});
test('agent run timeout fails without retrying', async () => {
const pool = createFakePool();
const gateway = createAgentRunGateway({
pool,
userAuth: {},
tkmindProxy: {
async startSessionForUser() {
return { id: 'session-timeout' };
},
async submitSessionReplyForUser() {
await new Promise(() => {});
},
},
retryDelaysMs: [0, 0],
maxConcurrentRuns: 1,
runTimeoutMs: 5,
});
const run = await gateway.createRun('user-1', {
requestId: 'req-timeout',
userMessage: { role: 'user', content: [] },
});
await waitFor(() => pool.runs.get(run.id)?.status === 'failed');
assert.equal(pool.runs.get(run.id).attempts, 1);
assert.match(pool.runs.get(run.id).error_message, /timed out/);
assert.equal(
pool.events.some((event) => event.runId === run.id && event.eventType === 'timeout'),
true,
);
});
test('agent run queue status reports active database and local queue state', async () => {
const pool = createFakePool();
const gateway = createAgentRunGateway({
pool,
userAuth: {},
tkmindProxy: {},
autoDispatch: false,
maxConcurrentRuns: 2,
runTimeoutMs: 1234,
});
await gateway.createRun('user-1', {
requestId: 'req-status',
userMessage: { role: 'user', content: [] },
});
const status = await gateway.getQueueStatus();
assert.equal(status.maxConcurrentRuns, 2);
assert.equal(status.runTimeoutMs, 1234);
assert.equal(status.inFlight, 0);
assert.equal(status.pendingDispatches, 0);
assert.equal(status.statusCounts.queued, 1);
});
+1
View File
@@ -346,6 +346,7 @@ async function writeMetadata() {
' node scripts/runtime-worker-drain.mjs undrain goosed-3',
' node scripts/check-tool-runtime.mjs',
' node scripts/check-agent-code-run-entry.mjs',
' curl -sk https://mm.tkmind.cn/api/runtime/status # includes toolRuntime.queue',
'',
].join('\n'),
);
+5
View File
@@ -148,6 +148,7 @@ function summarizeRuntime(runtimeJson) {
routerEnabled: Boolean(runtimeJson?.router?.enabled),
publicBaseUrl: runtimeJson?.publicBaseUrl ?? null,
toolRuntime: runtimeJson?.toolRuntime ?? null,
toolQueue: runtimeJson?.toolRuntime?.queue ?? null,
workers: workers.map((worker) => ({
id: worker.id,
healthy: runtimeJson?.targets?.find((target) => target.target === worker.target)?.healthy ?? null,
@@ -220,6 +221,10 @@ for (const worker of runtimeSummary?.workers ?? []) {
if (worker.firstTokenCount > 0 && worker.ewmaFirstTokenMs <= 0) failures.push(`${worker.id}_first_token_ewma_missing`);
}
if (runtimeSummary?.toolRuntime?.chatInjectsCodeTools !== false) failures.push('chat_injects_code_tools');
if (runtimeSummary?.toolQueue?.error) failures.push('tool_queue_status_error');
if (runtimeSummary?.toolQueue?.inFlight > runtimeSummary?.toolQueue?.maxConcurrentRuns) {
failures.push('tool_queue_concurrency_exceeded');
}
const report = {
ok: failures.length === 0,
+13
View File
@@ -572,6 +572,8 @@ async function bootstrapUserAuth() {
pool,
userAuth,
tkmindProxy,
maxConcurrentRuns: Number(process.env.MEMIND_AGENT_RUN_QUEUE_CONCURRENCY ?? 1),
runTimeoutMs: Number(process.env.MEMIND_AGENT_RUN_TIMEOUT_MS ?? 15 * 60 * 1000),
});
wechatMpService = createWechatMpService({
config: WECHAT_MP_CONFIG,
@@ -1781,6 +1783,17 @@ api.get('/runtime/status', async (_req, res) => {
}
try {
const status = await tkmindProxy.getRuntimeStatus();
const toolQueue = agentRunGateway?.getQueueStatus
? await agentRunGateway.getQueueStatus().catch((err) => ({
error: err instanceof Error ? err.message : String(err),
}))
: null;
if (toolQueue) {
status.toolRuntime = {
...(status.toolRuntime ?? {}),
queue: toolQueue,
};
}
return res.json({
ok: true,
timestamp: new Date().toISOString(),