memind/services/orchestrator/executor-worker-protocol.mjs

import crypto from 'node:crypto';
import {
  buildExecutorJobEvent,
  executorGatewayInternals,
} from './executor-gateway.mjs';

const WORKER_PROTOCOL_VERSION = 'executor-worker-v1';
const DEFAULT_LEASE_DURATION_MS = 30_000;
const DEFAULT_MAX_ATTEMPTS = 3;
const MAX_PROGRESS_DATA_BYTES = 16 * 1024;
const ACTIVE_STATUSES = Object.freeze(['leased', 'running']);

function clampInteger(value, fallback, min, max) {
  const number = Number(value);
  if (!Number.isFinite(number)) return fallback;
  return Math.min(max, Math.max(min, Math.floor(number)));
}

function protocolError(code, message, status = 409, details = null) {
  const error = new Error(message);
  error.code = code;
  error.status = status;
  error.details = details;
  return error;
}

function requiredIdentifier(value, name, max = 128) {
  const normalized = String(value ?? '').trim();
  if (!normalized || !/^[a-zA-Z0-9:_-]+$/.test(normalized) || normalized.length > max) {
    throw protocolError(
      `EXECUTOR_${name.toUpperCase()}_INVALID`,
      `${name} must be a safe identifier no longer than ${max} characters`,
      422,
    );
  }
  return normalized;
}

function sanitizeData(value, maxBytes = MAX_PROGRESS_DATA_BYTES) {
  if (value == null) return null;
  const encoded = JSON.stringify(value);
  if (Buffer.byteLength(encoded, 'utf8') > maxBytes) {
    throw protocolError(
      'EXECUTOR_EVENT_DATA_TOO_LARGE',
      `Executor event data exceeds ${maxBytes} bytes`,
      413,
    );
  }
  return structuredClone(value);
}

function sanitizeError(error) {
  const source = error && typeof error === 'object' ? error : {};
  return {
    code: String(source.code ?? 'EXECUTOR_FAILED').trim().slice(0, 128),
    message: String(source.message ?? error ?? 'Executor failed').trim().slice(0, 2048),
    retryable: source.retryable === true,
  };
}

function sanitizeResult(result) {
  const source = result && typeof result === 'object' ? result : {};
  const artifactRefs = Array.isArray(source.artifactRefs)
    ? source.artifactRefs
      .map((reference) => {
        const kind = String(reference?.kind ?? '').trim().slice(0, 64);
        const id = String(reference?.id ?? '').trim().slice(0, 256);
        return kind && id ? { kind, id } : null;
      })
      .filter(Boolean)
      .slice(0, 100)
    : [];
  return {
    outcome: String(source.outcome ?? 'completed').trim().slice(0, 128) || 'completed',
    summary: String(source.summary ?? '').trim().slice(0, 4096) || null,
    artifactRefs,
    metrics: source.metrics && typeof source.metrics === 'object'
      ? sanitizeData(source.metrics, 4096)
      : {},
  };
}

function requireWorkerStore(store) {
  for (const method of [
    'claimNext',
    'transition',
    'listExpiredLeases',
    'getQueueStats',
  ]) {
    if (typeof store?.[method] !== 'function') {
      throw new Error(`Executor Worker store missing method: ${method}`);
    }
  }
}

function requireApplied(result, action) {
  if (result?.applied) return result.record;
  const code = result?.reason === 'not_found'
    ? 'EXECUTOR_JOB_NOT_FOUND'
    : result?.reason === 'lease_mismatch'
      ? 'EXECUTOR_LEASE_LOST'
      : 'EXECUTOR_JOB_STATE_CONFLICT';
  const status = result?.reason === 'not_found' ? 404 : 409;
  throw protocolError(code, `Cannot ${action} executor job: ${result?.reason ?? 'conflict'}`, status, {
    currentStatus: result?.record?.status ?? null,
  });
}

export function createExecutorWorkerCoordinator({
  store,
  nowMs = Date.now,
  randomUUID = crypto.randomUUID,
  defaultLeaseDurationMs = DEFAULT_LEASE_DURATION_MS,
  defaultMaxAttempts = DEFAULT_MAX_ATTEMPTS,
  workerStaleAfterMs = 60_000,
} = {}) {
  requireWorkerStore(store);

  async function claim({
    workerId,
    executors,
    leaseDurationMs = defaultLeaseDurationMs,
    draining = false,
    adapterHealth = {},
  } = {}) {
    const normalizedWorkerId = requiredIdentifier(workerId, 'worker_id');
    const normalizedExecutors = [...new Set(
      (Array.isArray(executors) ? executors : [executors])
        .map((executor) => String(executor ?? '').trim().toLowerCase())
        .filter(Boolean),
    )].slice(0, 20);
    if (!normalizedExecutors.length) {
      throw protocolError(
        'EXECUTOR_WORKER_EXECUTORS_REQUIRED',
        'Worker claim requires at least one executor',
        422,
      );
    }
    const now = nowMs();
    await store.recordWorkerHeartbeat?.({
      version: WORKER_PROTOCOL_VERSION,
      workerId: normalizedWorkerId,
      executors: normalizedExecutors,
      draining: draining === true,
      adapterHealth: Object.fromEntries(
        normalizedExecutors.map((executor) => [
          executor,
          {
            ok: adapterHealth?.[executor]?.ok === true,
            code: adapterHealth?.[executor]?.code == null
              ? null
              : String(adapterHealth[executor].code).slice(0, 128),
          },
        ]),
      ),
      lastSeenAt: now,
    });
    if (draining === true) return null;
    const duration = clampInteger(leaseDurationMs, defaultLeaseDurationMs, 5_000, 5 * 60_000);
    const leaseToken = randomUUID();
    const claimed = await store.claimNext({
      workerId: normalizedWorkerId,
      executors: normalizedExecutors,
      now,
      leaseToken,
      leaseDurationMs: duration,
      eventFactory: (updated) => buildExecutorJobEvent({
        jobId: updated.jobId,
        type: 'executor_job_leased',
        timestamp: now,
        data: {
          workerId: normalizedWorkerId,
          executor: updated.executor,
          attempt: updated.attempts,
          leaseExpiresAt: updated.lease.expiresAt,
        },
      }),
    });
    return claimed
      ? {
          version: WORKER_PROTOCOL_VERSION,
          job: claimed,
          leaseToken,
        }
      : null;
  }

  async function start(jobId, { leaseToken } = {}) {
    const normalizedJobId = requiredIdentifier(jobId, 'job_id');
    const normalizedLeaseToken = requiredIdentifier(leaseToken, 'lease_token');
    const now = nowMs();
    const result = await store.transition(normalizedJobId, {
      expectedStatuses: ['leased'],
      leaseToken: normalizedLeaseToken,
      updater(current) {
        return {
          ...current,
          status: 'running',
          startedAt: current.startedAt ?? now,
          updatedAt: now,
        };
      },
      eventFactory: (updated) => buildExecutorJobEvent({
        jobId: updated.jobId,
        type: 'executor_job_started',
        timestamp: now,
        data: {
          workerId: updated.lease.workerId,
          attempt: updated.attempts,
        },
      }),
    });
    return requireApplied(result, 'start');
  }

  async function heartbeat(jobId, {
    leaseToken,
    leaseDurationMs = defaultLeaseDurationMs,
  } = {}) {
    const normalizedJobId = requiredIdentifier(jobId, 'job_id');
    const normalizedLeaseToken = requiredIdentifier(leaseToken, 'lease_token');
    const now = nowMs();
    const duration = clampInteger(leaseDurationMs, defaultLeaseDurationMs, 5_000, 5 * 60_000);
    const result = await store.transition(normalizedJobId, {
      expectedStatuses: [...ACTIVE_STATUSES],
      leaseToken: normalizedLeaseToken,
      updater(current) {
        return {
          ...current,
          lease: {
            ...current.lease,
            heartbeatAt: now,
            expiresAt: now + duration,
          },
          updatedAt: now,
        };
      },
      eventFactory: (updated) => buildExecutorJobEvent({
        jobId: updated.jobId,
        type: 'executor_job_heartbeat',
        timestamp: now,
        data: {
          workerId: updated.lease.workerId,
          leaseExpiresAt: updated.lease.expiresAt,
        },
      }),
    });
    return requireApplied(result, 'heartbeat');
  }

  async function progress(jobId, {
    leaseToken,
    type = 'executor_job_progress',
    data = null,
  } = {}) {
    const normalizedJobId = requiredIdentifier(jobId, 'job_id');
    const normalizedLeaseToken = requiredIdentifier(leaseToken, 'lease_token');
    const normalizedType = String(type ?? '').trim();
    if (!/^executor_job_[a-z0-9_]{1,96}$/.test(normalizedType)) {
      throw protocolError(
        'EXECUTOR_PROGRESS_TYPE_INVALID',
        'Progress event type must use the executor_job_ namespace',
        422,
      );
    }
    const safeData = sanitizeData(data);
    const now = nowMs();
    const result = await store.transition(normalizedJobId, {
      expectedStatuses: [...ACTIVE_STATUSES],
      leaseToken: normalizedLeaseToken,
      updater(current) {
        return { ...current, updatedAt: now };
      },
      eventFactory: (updated) => buildExecutorJobEvent({
        jobId: updated.jobId,
        type: normalizedType,
        timestamp: now,
        data: safeData,
      }),
    });
    return requireApplied(result, 'append progress');
  }

  async function complete(jobId, { leaseToken, result: output = null } = {}) {
    const normalizedJobId = requiredIdentifier(jobId, 'job_id');
    const normalizedLeaseToken = requiredIdentifier(leaseToken, 'lease_token');
    const safeResult = sanitizeResult(output);
    const now = nowMs();
    const transition = await store.transition(normalizedJobId, {
      expectedStatuses: [...ACTIVE_STATUSES],
      leaseToken: normalizedLeaseToken,
      updater(current) {
        return {
          ...current,
          status: 'succeeded',
          reason: null,
          result: safeResult,
          lease: null,
          updatedAt: now,
          completedAt: now,
        };
      },
      eventFactory: (updated) => buildExecutorJobEvent({
        jobId: updated.jobId,
        type: 'executor_job_succeeded',
        timestamp: now,
        data: safeResult,
      }),
    });
    return requireApplied(transition, 'complete');
  }

  async function fail(jobId, {
    leaseToken,
    error,
    retryDelayMs = 1_000,
  } = {}) {
    const normalizedJobId = requiredIdentifier(jobId, 'job_id');
    const normalizedLeaseToken = requiredIdentifier(leaseToken, 'lease_token');
    const safeError = sanitizeError(error);
    const now = nowMs();
    const transition = await store.transition(normalizedJobId, {
      expectedStatuses: [...ACTIVE_STATUSES],
      leaseToken: normalizedLeaseToken,
      updater(current) {
        const maxAttempts = clampInteger(
          current.maxAttempts,
          defaultMaxAttempts,
          1,
          20,
        );
        const retryable = safeError.retryable && Number(current.attempts ?? 0) < maxAttempts;
        return {
          ...current,
          status: retryable ? 'retryable' : 'failed',
          reason: safeError.code,
          error: safeError,
          lease: null,
          nextAttemptAt: retryable
            ? now + clampInteger(retryDelayMs, 1_000, 0, 60 * 60_000)
            : null,
          updatedAt: now,
          completedAt: retryable ? null : now,
        };
      },
      eventFactory: (updated) => buildExecutorJobEvent({
        jobId: updated.jobId,
        type: updated.status === 'retryable'
          ? 'executor_job_retry_scheduled'
          : 'executor_job_failed',
        timestamp: now,
        data: {
          status: updated.status,
          error: safeError,
          nextAttemptAt: updated.nextAttemptAt,
        },
      }),
    });
    return requireApplied(transition, 'fail');
  }

  async function recoverExpired({ limit = 100 } = {}) {
    const now = nowMs();
    const expired = await store.listExpiredLeases({
      now,
      limit: clampInteger(limit, 100, 1, 500),
    });
    const recovered = [];
    for (const candidate of expired) {
      const leaseToken = candidate.lease?.token;
      if (!leaseToken) continue;
      const transition = await store.transition(candidate.jobId, {
        expectedStatuses: [...ACTIVE_STATUSES],
        leaseToken,
        updater(current) {
          const maxAttempts = clampInteger(
            current.maxAttempts,
            defaultMaxAttempts,
            1,
            20,
          );
          const retryable = Number(current.attempts ?? 0) < maxAttempts;
          return {
            ...current,
            status: retryable ? 'retryable' : 'timed_out',
            reason: 'executor_lease_expired',
            error: {
              code: 'EXECUTOR_LEASE_EXPIRED',
              message: 'Worker lease expired before completion',
              retryable,
            },
            lease: null,
            nextAttemptAt: retryable ? now : null,
            updatedAt: now,
            completedAt: retryable ? null : now,
          };
        },
        eventFactory: (updated) => buildExecutorJobEvent({
          jobId: updated.jobId,
          type: updated.status === 'retryable'
            ? 'executor_job_lease_recovered'
            : 'executor_job_timed_out',
          timestamp: now,
          data: {
            previousWorkerId: candidate.lease?.workerId ?? null,
            status: updated.status,
            attempt: updated.attempts,
          },
        }),
      });
      if (transition.applied) recovered.push(transition.record);
    }
    return {
      version: WORKER_PROTOCOL_VERSION,
      inspected: expired.length,
      recovered,
    };
  }

  async function stats() {
    const now = nowMs();
    const workers = typeof store.listWorkers === 'function'
      ? await store.listWorkers({ now, staleAfterMs: workerStaleAfterMs })
      : [];
    return {
      version: WORKER_PROTOCOL_VERSION,
      ...(await store.getQueueStats({ now })),
      workers,
      healthyWorkers: workers.filter((worker) => (
        !worker.stale
        && !worker.draining
        && Object.values(worker.adapterHealth ?? {}).some((health) => health?.ok === true)
      )).length,
    };
  }

  return {
    claim,
    start,
    heartbeat,
    progress,
    complete,
    fail,
    recoverExpired,
    stats,
  };
}

export const executorWorkerProtocolInternals = {
  ACTIVE_STATUSES,
  DEFAULT_LEASE_DURATION_MS,
  DEFAULT_MAX_ATTEMPTS,
  MAX_PROGRESS_DATA_BYTES,
  WORKER_PROTOCOL_VERSION,
  clampInteger,
  protocolError,
  requiredIdentifier,
  sanitizeData,
  sanitizeError,
  sanitizeResult,
  terminalStatuses: executorGatewayInternals.TERMINAL_JOB_STATUSES,
};