memind/agent-run-gateway.mjs

import crypto from 'node:crypto';

const DEFAULT_RUN_RETRY_DELAYS_MS = [1500, 5000, 15000];
const TERMINAL_STATUSES = new Set(['succeeded', 'failed']);

function nowMs() {
  return Date.now();
}

function safeJsonParse(value, fallback = null) {
  try {
    return JSON.parse(value);
  } catch {
    return fallback;
  }
}

function serializeMessage(message) {
  return JSON.stringify(message ?? {});
}

function projectRun(row) {
  if (!row) return null;
  return {
    id: row.id,
    userId: row.user_id,
    sessionId: row.agent_session_id ?? null,
    requestId: row.request_id,
    status: row.status,
    attempts: Number(row.attempts ?? 0),
    error: row.error_message ?? null,
    createdAt: Number(row.created_at ?? 0),
    updatedAt: Number(row.updated_at ?? 0),
    startedAt: row.started_at == null ? null : Number(row.started_at),
    completedAt: row.completed_at == null ? null : Number(row.completed_at),
  };
}

export function createAgentRunGateway({
  pool,
  userAuth,
  tkmindProxy,
  retryDelaysMs = DEFAULT_RUN_RETRY_DELAYS_MS,
  autoDispatch = true,
}) {
  const inFlight = new Set();

  async function appendEvent(runId, type, data = null) {
    await pool.query(
      `INSERT INTO h5_agent_run_events (id, run_id, event_type, data_json, created_at)
       VALUES (?, ?, ?, ?, ?)`,
      [
        crypto.randomUUID(),
        runId,
        type,
        data == null ? null : JSON.stringify(data),
        nowMs(),
      ],
    );
  }

  async function getRunById(runId) {
    const [rows] = await pool.query(
      `SELECT * FROM h5_agent_runs WHERE id = ? LIMIT 1`,
      [runId],
    );
    return rows[0] ?? null;
  }

  async function getRunForUser(userId, runId) {
    const [rows] = await pool.query(
      `SELECT * FROM h5_agent_runs WHERE id = ? AND user_id = ? LIMIT 1`,
      [runId, userId],
    );
    return projectRun(rows[0] ?? null);
  }

  async function getRunByRequest(userId, requestId) {
    const [rows] = await pool.query(
      `SELECT * FROM h5_agent_runs WHERE user_id = ? AND request_id = ? LIMIT 1`,
      [userId, requestId],
    );
    return rows[0] ?? null;
  }

  async function createRun(userId, { sessionId = null, requestId, userMessage }) {
    const normalizedRequestId = String(requestId ?? '').trim();
    if (!normalizedRequestId) {
      throw new Error('缺少 request_id');
    }
    const existing = await getRunByRequest(userId, normalizedRequestId);
    if (existing) {
      if (autoDispatch && !TERMINAL_STATUSES.has(existing.status)) dispatchRun(existing.id);
      return projectRun(existing);
    }

    const runId = crypto.randomUUID();
    const createdAt = nowMs();
    await pool.query(
      `INSERT INTO h5_agent_runs
         (id, user_id, agent_session_id, request_id, status, attempts,
          user_message_json, error_message, created_at, updated_at, started_at, completed_at)
       VALUES (?, ?, ?, ?, 'queued', 0, ?, NULL, ?, ?, NULL, NULL)`,
      [
        runId,
        userId,
        sessionId || null,
        normalizedRequestId,
        serializeMessage(userMessage),
        createdAt,
        createdAt,
      ],
    );
    await appendEvent(runId, 'queued', { sessionId: sessionId || null });
    if (autoDispatch) dispatchRun(runId);
    return projectRun(await getRunById(runId));
  }

  async function markRun(runId, status, fields = {}) {
    const updates = ['status = ?', 'updated_at = ?'];
    const values = [status, nowMs()];
    for (const [key, value] of Object.entries(fields)) {
      updates.push(`${key} = ?`);
      values.push(value);
    }
    values.push(runId);
    await pool.query(
      `UPDATE h5_agent_runs SET ${updates.join(', ')} WHERE id = ?`,
      values,
    );
    await appendEvent(runId, status, fields);
  }

  async function processRun(runId) {
    const row = await getRunById(runId);
    if (!row || TERMINAL_STATUSES.has(row.status)) return;
    const nextAttempt = Number(row.attempts ?? 0) + 1;
    const [claim] = await pool.query(
      `UPDATE h5_agent_runs
       SET status = 'running', attempts = ?, started_at = COALESCE(started_at, ?), updated_at = ?, error_message = NULL
       WHERE id = ? AND status IN ('queued', 'retryable')`,
      [nextAttempt, nowMs(), nowMs(), runId],
    );
    if (Number(claim?.affectedRows ?? 0) === 0) return;
    await appendEvent(runId, 'running', { attempt: nextAttempt });

    try {
      let sessionId = row.agent_session_id ?? null;
      if (!sessionId) {
        const session = await tkmindProxy.startSessionForUser(row.user_id);
        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 });
      }

      const userMessage = safeJsonParse(row.user_message_json, {});
      await tkmindProxy.submitSessionReplyForUser(
        row.user_id,
        sessionId,
        row.request_id,
        userMessage,
      );
      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;
      await markRun(runId, retryable ? 'retryable' : 'failed', {
        error_message: message,
        completed_at: retryable ? null : nowMs(),
      });
      if (retryable) {
        setTimeout(() => dispatchRun(runId), retryDelaysMs[nextAttempt - 1]);
      }
    }
  }

  function dispatchRun(runId) {
    if (!runId || inFlight.has(runId)) return;
    inFlight.add(runId);
    void processRun(runId).finally(() => inFlight.delete(runId));
  }

  return {
    createRun,
    getRunForUser,
    dispatchRun,
  };
}