import assert from 'node:assert/strict'; import fs from 'node:fs/promises'; import os from 'node:os'; import path from 'node:path'; import test from 'node:test'; import { createAgentRunGateway } from './agent-run-gateway.mjs'; function createFakePool() { const runs = new Map(); const events = []; const latestHeartbeatAt = (runId) => { const timestamps = events .filter((event) => event.runId === runId && event.eventType === 'worker_heartbeat') .map((event) => Number(event.createdAt ?? 0)); return timestamps.length ? Math.max(...timestamps) : null; }; return { runs, events, async query(sql, params = []) { if (sql.includes('SELECT * FROM h5_agent_runs WHERE id = ? AND user_id = ?')) { const [id, userId] = params; const row = runs.get(id); return [[row && row.user_id === userId ? row : undefined].filter(Boolean)]; } if (sql.includes('SELECT * FROM h5_agent_runs WHERE user_id = ? AND request_id = ?')) { const [userId, requestId] = params; return [[...runs.values()].filter((row) => row.user_id === userId && row.request_id === requestId)]; } 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('h.latest_heartbeat_at') && sql.includes("WHERE r.status = 'running'") && sql.includes('LIMIT 1')) { return [[...runs.values()] .filter((row) => row.status === 'running') .map((row) => ({ id: row.id, request_id: row.request_id, started_at: row.started_at, updated_at: row.updated_at, attempts: row.attempts, latest_heartbeat_at: latestHeartbeatAt(row.id), })) .sort((a, b) => Number(a.latest_heartbeat_at ?? a.started_at ?? 0) - Number(b.latest_heartbeat_at ?? b.started_at ?? 0)) .slice(0, 1)]; } if (sql.includes('SELECT COUNT(*) AS count') && sql.includes('worker_heartbeat')) { return [[{ count: [...runs.values()] .filter((row) => row.status === 'running' && latestHeartbeatAt(row.id) == null) .length, }]]; } if (sql.includes('SELECT') && sql.includes('latest_heartbeat_at') && sql.includes('COALESCE(h.latest_heartbeat_at, r.started_at) <= ?')) { const [cutoff, limit = 1] = params; return [[...runs.values()] .map((row) => ({ row, heartbeatAt: latestHeartbeatAt(row.id) })) .filter(({ row, heartbeatAt }) => ( row.status === 'running' && row.started_at != null && Number(heartbeatAt ?? row.started_at) <= Number(cutoff) )) .sort((a, b) => Number(a.heartbeatAt ?? a.row.started_at) - Number(b.heartbeatAt ?? b.row.started_at)) .slice(0, Number(limit)) .map(({ row, heartbeatAt }) => ({ id: row.id, request_id: row.request_id, started_at: row.started_at, updated_at: row.updated_at, attempts: row.attempts, latest_heartbeat_at: heartbeatAt, }))]; } if (sql.includes('SELECT id') && sql.includes("status IN ('queued', 'retryable')")) { const limit = Number(params[0] ?? 1); return [[...runs.values()] .filter((row) => ['queued', 'retryable'].includes(row.status)) .sort((a, b) => Number(a.updated_at) - Number(b.updated_at)) .slice(0, limit) .map((row) => ({ id: row.id }))]; } if (sql.includes('INSERT INTO h5_agent_runs')) { const [ id, userId, sessionId, requestId, userMessageJson, createdAt, updatedAt, ] = params; runs.set(id, { id, user_id: userId, agent_session_id: sessionId, request_id: requestId, status: 'queued', attempts: 0, user_message_json: userMessageJson, error_message: null, created_at: createdAt, updated_at: updatedAt, started_at: null, completed_at: null, }); return [{ affectedRows: 1 }]; } if (sql.includes('INSERT INTO h5_agent_run_events')) { const [id, runId, eventType, dataJson, createdAt] = params; events.push({ id, runId, eventType, dataJson, createdAt }); return [{ affectedRows: 1 }]; } if (sql.includes("SET status = 'running'")) { const [attempts, startedAt, updatedAt, id] = params; const row = runs.get(id); if (!row || !['queued', 'retryable'].includes(row.status)) { return [{ affectedRows: 0 }]; } Object.assign(row, { status: 'running', attempts, started_at: row.started_at ?? startedAt, updated_at: updatedAt, error_message: null, }); return [{ affectedRows: 1 }]; } if (sql.includes("WHERE id = ?") && sql.includes("status = 'running'") && sql.includes('worker_heartbeat')) { const [errorMessage, updatedAt, completedAt, id, cutoff] = params; const row = runs.get(id); if ( !row || row.status !== 'running' || row.started_at == null || Number(latestHeartbeatAt(id) ?? row.started_at) > Number(cutoff) ) { return [{ affectedRows: 0 }]; } Object.assign(row, { status: 'failed', error_message: errorMessage, updated_at: updatedAt, completed_at: completedAt, }); return [{ affectedRows: 1 }]; } if (sql.includes('UPDATE h5_agent_runs SET')) { const id = params.at(-1); const row = runs.get(id); if (!row) return [{ affectedRows: 0 }]; const columns = [...sql.matchAll(/([a-z_]+) = \?/g)].map((match) => match[1]); for (let i = 0; i < columns.length; i += 1) { row[columns[i]] = params[i]; } return [{ affectedRows: 1 }]; } throw new Error(`Unhandled SQL: ${sql}`); }, }; } async function waitFor(predicate) { for (let i = 0; i < 50; i += 1) { if (predicate()) return; await new Promise((resolve) => setTimeout(resolve, 5)); } assert.fail('condition was not met'); } test('agent run creation is idempotent by user and request id', async () => { const pool = createFakePool(); const gateway = createAgentRunGateway({ pool, userAuth: {}, tkmindProxy: {}, autoDispatch: false, }); const first = await gateway.createRun('user-1', { requestId: 'req-1', userMessage: { role: 'user', content: [] }, }); const second = await gateway.createRun('user-1', { requestId: 'req-1', userMessage: { role: 'user', content: [] }, }); assert.equal(first.id, second.id); assert.equal(pool.runs.size, 1); assert.equal(second.status, 'queued'); }); test('agent run creation stores code tool metadata in the queued message', async () => { const pool = createFakePool(); const gateway = createAgentRunGateway({ pool, userAuth: {}, tkmindProxy: {}, autoDispatch: false, }); const run = await gateway.createRun('user-1', { requestId: 'req-code', userMessage: { role: 'user', content: [{ type: 'text', text: '修改仓库代码' }], metadata: { source: 'h5' }, }, toolMode: 'code-task', taskType: 'repo_refactor', }); const stored = JSON.parse(pool.runs.get(run.id).user_message_json); assert.equal(stored.metadata.source, 'h5'); assert.deepEqual(stored.metadata.memindRun, { toolMode: 'code', taskType: 'repo_refactor', }); }); test('agent run starts a session and marks submitted reply as succeeded', async () => { const pool = createFakePool(); const submitted = []; const gateway = createAgentRunGateway({ pool, userAuth: {}, tkmindProxy: { async startSessionForUser() { return { id: 'session-1' }; }, async submitSessionReplyForUser(userId, sessionId, requestId, userMessage) { submitted.push({ userId, sessionId, requestId, userMessage }); }, }, retryDelaysMs: [], }); const run = await gateway.createRun('user-1', { requestId: 'req-1', userMessage: { role: 'user', content: [{ type: 'text', text: 'hi' }] }, }); await waitFor(() => pool.runs.get(run.id)?.status === 'succeeded'); assert.deepEqual(submitted.map((item) => item.sessionId), ['session-1']); assert.equal(pool.runs.get(run.id).attempts, 1); }); test('agent run with code tool mode starts and submits with code policy', async () => { const pool = createFakePool(); const submitted = []; const codePolicy = { extensionOverrides: { aider: { allowed: true } }, }; const gateway = createAgentRunGateway({ pool, userAuth: { async getCodeAgentSessionPolicy(userId) { assert.equal(userId, 'user-1'); return codePolicy; }, }, tkmindProxy: { async startSessionForUser(userId, options = {}) { assert.equal(userId, 'user-1'); assert.equal(options.sessionPolicy, codePolicy); return { id: 'session-code' }; }, async submitSessionReplyForUser(userId, sessionId, requestId, userMessage, options = {}) { submitted.push({ userId, sessionId, requestId, userMessage, options }); }, }, retryDelaysMs: [], }); const run = await gateway.createRun('user-1', { requestId: 'req-code', userMessage: { role: 'user', content: [{ type: 'text', text: 'run code task' }] }, toolMode: 'code', }); await waitFor(() => pool.runs.get(run.id)?.status === 'succeeded'); assert.deepEqual(submitted.map((item) => item.options), [{ toolMode: 'code' }]); assert.equal(submitted[0].userMessage.metadata.memindRun.toolMode, 'code'); }); test('agent run with enabled tool gateway dispatches code runs outside goose session', async () => { const pool = createFakePool(); const jobs = []; const gateway = createAgentRunGateway({ pool, userAuth: { async resolveWorkingDir(userId) { assert.equal(userId, 'user-1'); return '/tmp/memind-user-1'; }, }, tkmindProxy: { async startSessionForUser() { assert.fail('goose session should not start for external tool gateway run'); }, async submitSessionReplyForUser() { assert.fail('goose reply should not be submitted for external tool gateway run'); }, }, toolGateway: { getStatus() { return { enabled: true, protocol: 'agent-run-v1' }; }, async executeJob(job) { jobs.push(job); return { ok: true, dryRun: true, executor: 'aider' }; }, }, retryDelaysMs: [], }); const run = await gateway.createRun('user-1', { requestId: 'req-code-tool-gateway', userMessage: { role: 'user', content: [{ type: 'text', text: 'run code task' }] }, toolMode: 'code', taskType: 'small_patch', }); await waitFor(() => pool.runs.get(run.id)?.status === 'succeeded'); assert.equal(jobs.length, 1); assert.equal(jobs[0].cwd, '/tmp/memind-user-1'); assert.equal(jobs[0].taskType, 'small_patch'); assert.equal(pool.runs.get(run.id).agent_session_id, null); assert.equal( pool.events.some((event) => event.runId === run.id && event.eventType === 'tool_gateway_dispatch'), true, ); assert.equal( pool.events.some((event) => event.runId === run.id && event.eventType === 'tool_gateway_result'), true, ); }); test('agent run validates expected tool gateway artifacts before succeeding', async () => { const pool = createFakePool(); const workdir = await fs.mkdtemp(path.join(os.tmpdir(), 'memind-tool-validation-')); const gateway = createAgentRunGateway({ pool, userAuth: { async resolveWorkingDir() { return workdir; }, }, tkmindProxy: { async startSessionForUser() { assert.fail('goose session should not start for external tool gateway run'); }, async submitSessionReplyForUser() { assert.fail('goose reply should not be submitted for external tool gateway run'); }, }, toolGateway: { getStatus() { return { enabled: true, protocol: 'agent-run-v1' }; }, async executeJob() { await fs.writeFile( path.join(workdir, 'RESULT.md'), 'validated artifact from tool gateway\n', 'utf8', ); return { ok: true, dryRun: false, executor: 'aider', exitCode: 0, cwd: workdir, stdout: 'created RESULT.md', stderr: '', }; }, }, retryDelaysMs: [], }); const run = await gateway.createRun('user-1', { requestId: 'req-code-tool-validation', userMessage: { role: 'user', content: [{ type: 'text', text: 'create validation artifact' }], metadata: { memindRun: { validation: { expectedFile: { path: 'RESULT.md', contains: 'validated artifact', }, }, }, }, }, toolMode: 'code', taskType: 'small_patch', }); await waitFor(() => pool.runs.get(run.id)?.status === 'succeeded'); const validationEvent = pool.events.find( (event) => event.runId === run.id && event.eventType === 'tool_gateway_validation', ); assert.ok(validationEvent); assert.equal(JSON.parse(validationEvent.dataJson).expectedFiles[0].path, 'RESULT.md'); }); test('agent run fails non-retryably when tool gateway artifact validation fails', async () => { const pool = createFakePool(); const workdir = await fs.mkdtemp(path.join(os.tmpdir(), 'memind-tool-validation-missing-')); let attempts = 0; const gateway = createAgentRunGateway({ pool, userAuth: { async resolveWorkingDir() { return workdir; }, }, tkmindProxy: { async startSessionForUser() { assert.fail('goose session should not start for external tool gateway run'); }, async submitSessionReplyForUser() { assert.fail('goose reply should not be submitted for external tool gateway run'); }, }, toolGateway: { getStatus() { return { enabled: true, protocol: 'agent-run-v1' }; }, async executeJob() { attempts += 1; return { ok: true, dryRun: false, executor: 'aider', exitCode: 0, cwd: workdir, }; }, }, retryDelaysMs: [0, 0], }); const run = await gateway.createRun('user-1', { requestId: 'req-code-tool-validation-fail', userMessage: { role: 'user', content: [{ type: 'text', text: 'forget to create validation artifact' }], metadata: { memindRun: { validation: { expectedFile: 'MISSING.md', }, }, }, }, toolMode: 'code', taskType: 'small_patch', }); await waitFor(() => pool.runs.get(run.id)?.status === 'failed'); assert.equal(attempts, 1); assert.equal(pool.runs.get(run.id).attempts, 1); assert.match(pool.runs.get(run.id).error_message, /expected file not found/); assert.equal( pool.events.some((event) => event.runId === run.id && event.eventType === 'tool_gateway_validation_failed'), true, ); }); test('agent run retries transient failures and then becomes terminal', async () => { const pool = createFakePool(); const gateway = createAgentRunGateway({ pool, userAuth: {}, tkmindProxy: { async startSessionForUser() { return { id: 'session-1' }; }, async submitSessionReplyForUser() { throw new Error('upstream unavailable'); }, }, retryDelaysMs: [0, 0], }); const run = await gateway.createRun('user-1', { requestId: 'req-1', userMessage: { role: 'user', content: [] }, }); await waitFor(() => pool.runs.get(run.id)?.status === 'failed'); 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); }); test('external worker dispatches queued runs through the same queue controls', async () => { const pool = createFakePool(); const submitted = []; const gateway = createAgentRunGateway({ pool, userAuth: {}, tkmindProxy: { async startSessionForUser() { return { id: 'session-worker' }; }, async submitSessionReplyForUser(_userId, _sessionId, requestId) { submitted.push(requestId); }, }, autoDispatch: false, retryDelaysMs: [], maxConcurrentRuns: 1, }); const run = await gateway.createRun('user-1', { requestId: 'req-worker', userMessage: { role: 'user', content: [] }, }); assert.equal(pool.runs.get(run.id).status, 'queued'); const result = await gateway.dispatchQueuedRuns({ limit: 10 }); assert.equal(result.dispatched, 1); await waitFor(() => pool.events.some((event) => event.runId === run.id && event.eventType === 'worker_heartbeat')); await waitFor(() => pool.runs.get(run.id)?.status === 'succeeded'); assert.deepEqual(submitted, ['req-worker']); }); test('external worker does not dispatch more runs when local queue is full', async () => { const pool = createFakePool(); const release = []; const gateway = createAgentRunGateway({ pool, userAuth: {}, tkmindProxy: { async startSessionForUser() { return { id: `session-${release.length + 1}` }; }, async submitSessionReplyForUser() { await new Promise((resolve) => release.push(resolve)); }, }, autoDispatch: false, retryDelaysMs: [], maxConcurrentRuns: 1, }); const run1 = await gateway.createRun('user-1', { requestId: 'req-full-1', userMessage: { role: 'user', content: [] }, }); const run2 = await gateway.createRun('user-1', { requestId: 'req-full-2', userMessage: { role: 'user', content: [] }, }); try { const first = await gateway.dispatchQueuedRuns({ limit: 10 }); assert.equal(first.dispatched, 1); await waitFor(() => pool.runs.get(run1.id)?.status === 'running'); assert.equal((await gateway.getQueueStatus()).inFlight, 1); const second = await gateway.dispatchQueuedRuns({ limit: 10 }); assert.equal(second.dispatched, 0); assert.equal(pool.runs.get(run2.id).status, 'queued'); } finally { while (release.length > 0) release.shift()(); } await waitFor(() => pool.runs.get(run1.id)?.status === 'succeeded'); assert.equal(pool.runs.get(run2.id).status, 'queued'); }); test('stale running recovery dry-run reports rows without mutating them', async () => { const pool = createFakePool(); const gateway = createAgentRunGateway({ pool, userAuth: {}, tkmindProxy: {}, autoDispatch: false, runTimeoutMs: 1000, }); const run = await gateway.createRun('user-1', { requestId: 'req-stale-dry-run', userMessage: { role: 'user', content: [] }, }); Object.assign(pool.runs.get(run.id), { status: 'running', attempts: 1, started_at: Date.now() - 5000, updated_at: Date.now() - 5000, }); const result = await gateway.recoverStaleRunningRuns({ staleMs: 1000, dryRun: true }); assert.equal(result.considered, 1); assert.equal(result.recovered, 0); assert.equal(result.runs[0].id, run.id); assert.equal(pool.runs.get(run.id).status, 'running'); assert.equal( pool.events.some((event) => event.runId === run.id && event.eventType === 'stale_recovered'), false, ); }); test('stale running recovery marks old running rows failed with an event', async () => { const pool = createFakePool(); const gateway = createAgentRunGateway({ pool, userAuth: {}, tkmindProxy: {}, autoDispatch: false, runTimeoutMs: 1000, }); const run = await gateway.createRun('user-1', { requestId: 'req-stale-apply', userMessage: { role: 'user', content: [] }, }); Object.assign(pool.runs.get(run.id), { status: 'running', attempts: 1, started_at: Date.now() - 5000, updated_at: Date.now() - 5000, }); const result = await gateway.recoverStaleRunningRuns({ staleMs: 1000, dryRun: false }); assert.equal(result.considered, 1); assert.equal(result.recovered, 1); assert.equal(pool.runs.get(run.id).status, 'failed'); assert.match(pool.runs.get(run.id).error_message, /stale running state/); assert.equal( pool.events.some((event) => event.runId === run.id && event.eventType === 'stale_recovered'), true, ); }); test('stale running recovery ignores old running rows with fresh heartbeat', async () => { const pool = createFakePool(); const gateway = createAgentRunGateway({ pool, userAuth: {}, tkmindProxy: {}, autoDispatch: false, runTimeoutMs: 1000, }); const run = await gateway.createRun('user-1', { requestId: 'req-stale-heartbeat-fresh', userMessage: { role: 'user', content: [] }, }); Object.assign(pool.runs.get(run.id), { status: 'running', attempts: 1, started_at: Date.now() - 5000, updated_at: Date.now() - 5000, }); pool.events.push({ id: 'heartbeat-1', runId: run.id, eventType: 'worker_heartbeat', dataJson: JSON.stringify({ attempt: 1 }), createdAt: Date.now(), }); const result = await gateway.recoverStaleRunningRuns({ staleMs: 1000, dryRun: false }); assert.equal(result.considered, 0); assert.equal(result.recovered, 0); assert.equal(pool.runs.get(run.id).status, 'running'); }); test('queue status reports running heartbeat age and missing heartbeat count', async () => { const pool = createFakePool(); const gateway = createAgentRunGateway({ pool, userAuth: {}, tkmindProxy: {}, autoDispatch: false, runTimeoutMs: 1000, heartbeatMs: 250, }); const runWithHeartbeat = await gateway.createRun('user-1', { requestId: 'req-heartbeat-status-1', userMessage: { role: 'user', content: [] }, }); Object.assign(pool.runs.get(runWithHeartbeat.id), { status: 'running', attempts: 1, started_at: Date.now() - 5000, updated_at: Date.now() - 5000, }); pool.events.push({ id: 'heartbeat-status-1', runId: runWithHeartbeat.id, eventType: 'worker_heartbeat', dataJson: JSON.stringify({ attempt: 1 }), createdAt: Date.now() - 100, }); const runWithoutHeartbeat = await gateway.createRun('user-1', { requestId: 'req-heartbeat-status-2', userMessage: { role: 'user', content: [] }, }); Object.assign(pool.runs.get(runWithoutHeartbeat.id), { status: 'running', attempts: 1, started_at: Date.now() - 2000, updated_at: Date.now() - 2000, }); const status = await gateway.getQueueStatus(); assert.equal(status.heartbeatMs, 250); assert.equal(status.runningWithoutHeartbeatCount, 1); assert.equal(status.latestRunningRun.id, runWithoutHeartbeat.id); assert.equal(status.oldestRunningHeartbeatAt, null); assert.ok(status.oldestRunningHeartbeatAgeMs >= 1900); });