memind/agent-run-gateway.test.mjs

import assert from 'node:assert/strict';
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('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 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 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/);
});