feat: add agent run queue controls

2026-07-02 07:52:07 +08:00
parent 8e5094fcd9
commit 3ea973968e
9 changed files with 387 additions and 66 deletions
@@ -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,
  };
 }