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
@@ -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);
+});