memind/services/orchestrator/runtime.mjs

import { buildRunEvent, normalizeRunSpec } from './contracts.mjs';
import {
  CODE_RUN_SHADOW_WORKFLOW,
  createCodeRunShadowGraph,
} from './code-run-shadow-graph.mjs';
import {
  createExecutorGateway,
  createInMemoryExecutorJobStore,
  createPhase5ExecutorAdapterRegistry,
} from './executor-gateway.mjs';
import { createExecutorWorkerCoordinator } from './executor-worker-protocol.mjs';
import { createExecutorAdmissionPolicy } from './executor-admission-policy.mjs';

function graphConfig(runId) {
  return {
    configurable: {
      thread_id: runId,
      checkpoint_ns: '',
    },
  };
}

function projectSnapshot(runId, snapshot) {
  const values = snapshot?.values ?? {};
  if (!values.spec) return null;
  return {
    version: 'orchestrator-state-v1',
    runId,
    requestId: values.spec.requestId,
    workflow: values.spec.workflow,
    status: values.status ?? 'unknown',
    phase: values.phase ?? null,
    plan: values.plan ?? null,
    result: values.result ?? null,
    createdAt: values.events?.[0]?.timestamp ?? null,
    updatedAt: values.events?.at?.(-1)?.timestamp ?? null,
  };
}

export function createLangGraphOrchestratorRuntime({
  checkpointer,
  checkpointKind = 'unknown',
  durable = false,
  checkpointProbe = null,
  executorJobStore = null,
  executorJobProbe = null,
  executionEnabled = false,
  enabledExecutors = [],
  workerOptions = {},
  admissionConfig = {},
} = {}) {
  const jobStore = executorJobStore ?? createInMemoryExecutorJobStore();
  const admissionPolicy = createExecutorAdmissionPolicy({
    store: jobStore,
    config: admissionConfig,
    nowMs: workerOptions.nowMs,
  });
  const executorGateway = createExecutorGateway({
    store: jobStore,
    registry: createPhase5ExecutorAdapterRegistry({ enabledExecutors }),
    executionEnabled,
    admissionPolicy,
  });
  const workerCoordinator = createExecutorWorkerCoordinator({
    store: jobStore,
    ...workerOptions,
  });
  const graph = createCodeRunShadowGraph({
    checkpointer,
    executorGateway,
    executionEnabled,
  });
  const probe = typeof checkpointProbe === 'function'
    ? checkpointProbe
    : async () => {
        await checkpointer.getTuple(graphConfig('__orchestrator_readiness__'));
        return true;
      };
  const probeExecutorJobs = typeof executorJobProbe === 'function'
    ? executorJobProbe
    : async () => true;

  async function getSnapshot(runId) {
    return graph.getState(graphConfig(runId));
  }

  async function refreshExecutorTerminal(runId, snapshot) {
    const values = snapshot?.values ?? {};
    if (values.status !== 'waiting') return snapshot;
    const jobId = values.plan?.executorJob?.id;
    if (!jobId) return snapshot;
    const job = await executorGateway.getJob(jobId);
    const terminal = {
      succeeded: { status: 'succeeded', phase: 'completed' },
      failed: { status: 'failed', phase: 'executor_failed' },
      cancelled: { status: 'cancelled', phase: 'cancelled' },
      timed_out: { status: 'failed', phase: 'executor_timed_out' },
    }[job?.status];
    if (!terminal) return snapshot;
    const result = {
      observed: false,
      executed: job.status === 'succeeded',
      executorAdapter: values.plan.executorAdapter,
      taskType: values.plan.taskType,
      executorJob: {
        id: job.jobId,
        status: job.status,
        reason: job.reason,
        result: job.result ?? null,
      },
    };
    await graph.updateState(graphConfig(runId), {
      ...terminal,
      result,
      events: [buildRunEvent({
        runId,
        sequence: (values.events?.length ?? 0) + 1,
        type: job.status === 'succeeded'
          ? 'workflow_completed'
          : `workflow_${job.status}`,
        data: result,
      })],
    });
    return getSnapshot(runId);
  }

  async function getState(runId) {
    const normalizedRunId = String(runId ?? '').trim();
    if (!normalizedRunId) return null;
    const snapshot = await getSnapshot(normalizedRunId);
    return projectSnapshot(
      normalizedRunId,
      await refreshExecutorTerminal(normalizedRunId, snapshot),
    );
  }

  function health() {
    return {
      status: 'ok',
      service: 'memind-langgraph-orchestrator',
      checkpoint: {
        kind: checkpointKind,
        durable: Boolean(durable),
      },
      execution: executionEnabled ? 'worker-queue-enabled' : 'observe-only',
      executorGateway: executorGateway.status(),
      workerProtocol: {
        version: 'executor-worker-v1',
        enabled: executionEnabled === true,
      },
      admissionPolicy: admissionPolicy.status(),
    };
  }

  return {
    async start(input) {
      const spec = normalizeRunSpec(input);
      if (spec.workflow.name !== CODE_RUN_SHADOW_WORKFLOW) {
        const error = new Error(`Unsupported workflow: ${spec.workflow.name}`);
        error.code = 'WORKFLOW_NOT_SUPPORTED';
        error.status = 422;
        throw error;
      }
      const existing = await getState(spec.runId);
      if (existing) return existing;
      await graph.invoke({
        spec,
        status: 'queued',
        phase: 'accepted',
        events: [],
      }, graphConfig(spec.runId));
      return getState(spec.runId);
    },

    getState,

    async listEvents(runId, { after = 0 } = {}) {
      const normalizedRunId = String(runId ?? '').trim();
      const initial = normalizedRunId ? await getSnapshot(normalizedRunId) : null;
      const snapshot = normalizedRunId
        ? await refreshExecutorTerminal(normalizedRunId, initial)
        : null;
      if (!snapshot?.values?.spec) return null;
      const cursor = Math.max(0, Number(after) || 0);
      const events = (snapshot.values.events ?? []).filter((event) => event.sequence > cursor);
      return {
        runId: normalizedRunId,
        events,
        nextCursor: events.at(-1)?.sequence ?? cursor,
      };
    },

    async resume(runId) {
      const state = await getState(runId);
      if (!state) return null;
      if (state.status === 'waiting') return state;
      const error = new Error('This workflow has no interrupt point to resume');
      error.code = 'WORKFLOW_NOT_INTERRUPTED';
      error.status = 409;
      throw error;
    },

    async cancel(runId) {
      const state = await getState(runId);
      if (!state) return null;
      if (state.status === 'waiting' && state.plan?.executorJob?.id) {
        await executorGateway.cancel(state.plan.executorJob.id, {
          reason: 'workflow_cancelled',
        });
        return getState(runId);
      }
      const error = new Error('Completed shadow observations cannot be cancelled');
      error.code = 'WORKFLOW_ALREADY_TERMINAL';
      error.status = 409;
      throw error;
    },

    getExecutorJob(jobId) {
      return executorGateway.getJob(jobId);
    },

    listExecutorJobEvents(jobId, options = {}) {
      return executorGateway.listEvents(jobId, options);
    },

    submitExecutorJob(input) {
      return executorGateway.createJob(input);
    },

    cancelExecutorJob(jobId, options = {}) {
      return executorGateway.cancel(jobId, options);
    },

    claimExecutorJob(input) {
      return workerCoordinator.claim(input);
    },

    startExecutorJob(jobId, input) {
      return workerCoordinator.start(jobId, input);
    },

    heartbeatExecutorJob(jobId, input) {
      return workerCoordinator.heartbeat(jobId, input);
    },

    progressExecutorJob(jobId, input) {
      return workerCoordinator.progress(jobId, input);
    },

    completeExecutorJob(jobId, input) {
      return workerCoordinator.complete(jobId, input);
    },

    failExecutorJob(jobId, input) {
      return workerCoordinator.fail(jobId, input);
    },

    recoverExpiredExecutorJobs(input = {}) {
      return workerCoordinator.recoverExpired(input);
    },

    getExecutorQueueStats() {
      return workerCoordinator.stats();
    },

    async metrics() {
      return {
        version: 'orchestrator-metrics-v1',
        timestamp: Date.now(),
        queue: await workerCoordinator.stats(),
        gateway: executorGateway.status(),
      };
    },

    health,

    async ready() {
      await Promise.all([probe(), probeExecutorJobs()]);
      const executorQueue = await workerCoordinator.stats();
      if (executionEnabled && executorQueue.healthyWorkers < 1) {
        const error = new Error('Execution is enabled but no healthy executor worker is registered');
        error.code = 'EXECUTOR_WORKER_UNAVAILABLE';
        throw error;
      }
      return {
        ...health(),
        ready: true,
        executorQueue,
      };
    },
  };
}

export const orchestratorRuntimeInternals = {
  graphConfig,
  projectSnapshot,
};