Your Name
613 lines
19 KiB
JavaScript
613 lines
19 KiB
JavaScript
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 = [];
|
|
|
|
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('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('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.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');
|
|
});
|