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
+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(),