trustgraph/ts/packages/base/src/messaging/runtime.ts

/**
 * Effect-native messaging factories and scoped runtime helpers.
 */

import { randomUUID } from "node:crypto";
import {
  Context,
  Deferred,
  Duration,
  Effect,
  Fiber,
  Layer,
  PubSub as EffectPubSub,
  Ref,
  Result,
  Schedule,
  Scope,
  Stream,
} from "effect";
import * as O from "effect/Option";
import * as S from "effect/Schema";
import type {
  BackendConsumer,
  BackendProducer,
  CreateConsumerOptions,
  CreateProducerOptions,
  Message,
} from "../backend/types.js";
import { PubSub, type PubSubService } from "../backend/pubsub.js";
import {
  flowRuntimeError,
  messagingDeliveryError,
  messagingHandlerError,
  messagingLifecycleError,
  messagingTimeoutError,
  TooManyRequestsError,
  type FlowRuntimeError,
  type MessagingDeliveryError,
  type MessagingHandlerError,
  type MessagingLifecycleError,
  type MessagingTimeoutError,
  type PubSubError,
} from "../errors.js";
import type { ProducerMetrics } from "../metrics/index.js";
import type { FlowContext } from "./consumer.js";
import type { Flow } from "../processor/flow.js";
import type { SpecRuntimeRequirements } from "../spec/types.js";
import {
  loadMessagingRuntimeConfig,
  type MessagingRuntimeConfig,
} from "../runtime/messaging-config.js";

const isTooManyRequestsError = S.is(TooManyRequestsError);

export type EffectMessageHandler<T, E = never, R = never> = (
  message: T,
  properties: Record<string, string>,
  flow: FlowContext<R>,
) => Effect.Effect<void, E, R>;

export interface EffectProducerOptions<T = unknown> {
  readonly topic: string;
  readonly schema?: S.Codec<T, unknown>;
  readonly metrics?: ProducerMetrics;
}

export interface EffectProducer<T> {
  readonly send: (id: string, message: T) => Effect.Effect<void, MessagingDeliveryError>;
  readonly flush: Effect.Effect<void, MessagingDeliveryError>;
  readonly close: Effect.Effect<void, MessagingDeliveryError>;
}

export interface EffectConsumerOptions<T, E = never, R = never> {
  readonly topic: string;
  readonly subscription: string;
  readonly handler: EffectMessageHandler<T, E, R>;
  readonly concurrency?: number;
  readonly initialPosition?: "latest" | "earliest";
  readonly schema?: S.Codec<T, unknown>;
  readonly receiveTimeoutMs?: number;
  readonly errorBackoffMs?: number;
  readonly rateLimitRetryMs?: number;
  readonly rateLimitTimeoutMs?: number;
}

export interface EffectConsumer {
  readonly stop: Effect.Effect<void, MessagingLifecycleError>;
  readonly fibers: ReadonlyArray<Fiber.Fiber<void, never>>;
}

export interface EffectRequestResponseOptions<TReq = unknown, TRes = unknown> {
  readonly requestTopic: string;
  readonly responseTopic: string;
  readonly subscription: string;
  readonly requestSchema?: S.Codec<TReq, unknown>;
  readonly responseSchema?: S.Codec<TRes, unknown>;
}

export interface EffectRequestOptions<TRes, E = never, R = never> {
  readonly timeoutMs?: number;
  readonly recipient?: (response: TRes) => Effect.Effect<boolean, E, R>;
}

export interface EffectRequestResponse<TReq, TRes> {
  readonly request: <E = never, R = never>(
    request: TReq,
    options?: EffectRequestOptions<TRes, E, R>,
  ) => Effect.Effect<TRes, MessagingDeliveryError | MessagingLifecycleError | MessagingTimeoutError | E, R>;
  readonly stop: Effect.Effect<void, MessagingLifecycleError | MessagingDeliveryError>;
}

export interface ProducerFactoryService {
  readonly make: <T>(
    options: EffectProducerOptions<T>,
  ) => Effect.Effect<EffectProducer<T>, PubSubError, Scope.Scope>;
}

export interface ConsumerFactoryService {
  readonly run: <T, E = never, R = never>(
    options: EffectConsumerOptions<T, E, R>,
    flow: FlowContext<R>,
  ) => Effect.Effect<EffectConsumer, PubSubError, Scope.Scope | R>;
}

export interface RequestResponseFactoryService {
  readonly make: <TReq, TRes>(
    options: EffectRequestResponseOptions<TReq, TRes>,
  ) => Effect.Effect<EffectRequestResponse<TReq, TRes>, PubSubError, Scope.Scope>;
}

export interface FlowRuntimeService {
  readonly run: <Requirements = never>(
    flow: Flow<Requirements>,
  ) => Effect.Effect<void, FlowRuntimeError, SpecRuntimeRequirements | Requirements>;
}

interface ResponseEnvelope<T> {
  readonly id: string;
  readonly value: T;
}

export class ProducerFactory extends Context.Service<ProducerFactory, ProducerFactoryService>()(
  "@trustgraph/base/messaging/runtime/ProducerFactory",
) {}

export class ConsumerFactory extends Context.Service<ConsumerFactory, ConsumerFactoryService>()(
  "@trustgraph/base/messaging/runtime/ConsumerFactory",
) {}

export class RequestResponseFactory extends Context.Service<
  RequestResponseFactory,
  RequestResponseFactoryService
>()("@trustgraph/base/messaging/runtime/RequestResponseFactory") {}

export class FlowRuntime extends Context.Service<FlowRuntime, FlowRuntimeService>()(
  "@trustgraph/base/messaging/runtime/FlowRuntime",
) {}

export function makeEffectProducerHandle<T>(
  backend: BackendProducer<T>,
  options: EffectProducerOptions<T>,
): EffectProducer<T> {
  return {
    send: Effect.fn(`Producer.send:${options.topic}`)((id: string, message: T) =>
      Effect.tryPromise({
        try: () => backend.send(message, { id }),
        catch: (error) => messagingDeliveryError(options.topic, "send", error),
      }).pipe(
        Effect.tap(() =>
          options.metrics === undefined
            ? Effect.void
            : options.metrics.inc,
        ),
      ),
    ),
    flush: Effect.tryPromise({
      try: () => backend.flush(),
      catch: (error) => messagingDeliveryError(options.topic, "flush", error),
    }),
    close: Effect.tryPromise({
      try: () => backend.close(),
      catch: (error) => messagingDeliveryError(options.topic, "close", error),
    }),
  };
}

export const makeEffectProducerFromPubSub = Effect.fn("makeEffectProducerFromPubSub")(function* <T>(
  pubsub: PubSubService,
  options: EffectProducerOptions<T>,
) {
  const createOptions: CreateProducerOptions<T> = options.schema === undefined
    ? { topic: options.topic }
    : { topic: options.topic, schema: options.schema };
  const backend = yield* pubsub.createProducer<T>(createOptions);
  const producer = makeEffectProducerHandle(backend, options);

  yield* Effect.addFinalizer(() =>
    producer.close.pipe(
      Effect.catch((error) =>
        Effect.logError("[Producer] Failed to close producer", {
          error: error.message,
          topic: error.topic,
        }),
      ),
    ),
  );

  return producer;
});

const closeConsumerBackend = <T>(
  backend: BackendConsumer<T>,
  topic: string,
  subscription: string,
) =>
  Effect.tryPromise({
    try: () => backend.close(),
    catch: (error) => messagingLifecycleError(`${topic}:${subscription}`, "close-consumer", error),
  });

const acknowledgeMessage = <T>(
  backend: BackendConsumer<T>,
  message: Message<T>,
  topic: string,
) =>
  Effect.tryPromise({
    try: () => backend.acknowledge(message),
    catch: (error) => messagingDeliveryError(topic, "acknowledge", error),
  });

const negativeAcknowledgeMessage = <T>(
  backend: BackendConsumer<T>,
  message: Message<T>,
  topic: string,
) =>
  Effect.tryPromise({
    try: () => backend.negativeAcknowledge(message),
    catch: (error) => messagingDeliveryError(topic, "negative-acknowledge", error),
  });

const receiveMessage = <T>(
  backend: BackendConsumer<T>,
  topic: string,
  timeoutMs: number,
) =>
  Effect.tryPromise({
    try: () => backend.receive(timeoutMs),
    catch: (error) => messagingDeliveryError(topic, "receive", error),
  });

const handleMessageWithRetry = Effect.fn("handleMessageWithRetry")(function* <T, E, R>(
  options: EffectConsumerOptions<T, E, R>,
  flow: FlowContext<R>,
  message: Message<T>,
  config: MessagingRuntimeConfig,
) {
  const rateLimitRetryMs = options.rateLimitRetryMs ?? config.rateLimitRetryMs;
  const rateLimitTimeoutMs = options.rateLimitTimeoutMs ?? config.rateLimitTimeoutMs;
  const runHandler = (): Effect.Effect<void, TooManyRequestsError | MessagingHandlerError, R> =>
    options.handler(message.value(), message.properties(), flow).pipe(
      Effect.mapError((error): TooManyRequestsError | MessagingHandlerError =>
        isTooManyRequestsError(error)
          ? error
          : messagingHandlerError(options.topic, options.subscription, error),
      ),
    );

  return yield* runHandler().pipe(
    Effect.tapError((error) =>
      isTooManyRequestsError(error)
        ? Effect.logWarning("[Consumer] Rate limited, retrying", {
            topic: options.topic,
            subscription: options.subscription,
            retryMs: rateLimitRetryMs,
          })
        : Effect.void,
    ),
    Effect.retry({
      schedule: Schedule.spaced(Duration.millis(rateLimitRetryMs)),
      while: isTooManyRequestsError,
    }),
    Effect.timeoutOrElse({
      duration: Duration.millis(rateLimitTimeoutMs),
      orElse: () => Effect.fail(messagingTimeoutError("rate-limit", rateLimitTimeoutMs)),
    }),
    Effect.mapError((error) =>
      isTooManyRequestsError(error)
        ? messagingHandlerError(options.topic, options.subscription, error)
        : error,
    ),
  );
});

const processConsumerMessage = Effect.fn("processConsumerMessage")(function* <T, E, R>(
  backend: BackendConsumer<T>,
  options: EffectConsumerOptions<T, E, R>,
  flow: FlowContext<R>,
  message: Message<T>,
  config: MessagingRuntimeConfig,
) {
  yield* handleMessageWithRetry(options, flow, message, config).pipe(
    Effect.flatMap(() => acknowledgeMessage(backend, message, options.topic)),
    Effect.catch((error) =>
      negativeAcknowledgeMessage(backend, message, options.topic).pipe(
        Effect.catch((nakError) =>
          Effect.logError("[Consumer] Failed to negative-acknowledge message", {
            error: nakError.message,
            topic: nakError.topic,
          }),
        ),
        Effect.flatMap(() =>
          Effect.logError("[Consumer] Message handling failed", {
            error: error.message,
            topic: options.topic,
            subscription: options.subscription,
          }),
        ),
      ),
    ),
  );
});

const consumerLoop = <T, E, R>(
  backend: BackendConsumer<T>,
  options: EffectConsumerOptions<T, E, R>,
  flow: FlowContext<R>,
  config: MessagingRuntimeConfig,
): Effect.Effect<void, never, R> =>
  Effect.whileLoop({
    while: () => true,
    body: () =>
      receiveMessage(backend, options.topic, options.receiveTimeoutMs ?? config.consumerReceiveTimeoutMs).pipe(
        Effect.flatMap((message) =>
          message === null
            ? Effect.sleep(Duration.millis(options.receiveTimeoutMs ?? config.consumerReceiveTimeoutMs))
            : processConsumerMessage(backend, options, flow, message, config),
        ),
        Effect.catch((error) =>
          Effect.logError("[Consumer] Receive loop failed", {
            error: error.message,
            topic: options.topic,
            subscription: options.subscription,
          }).pipe(
            Effect.flatMap(() =>
              Effect.sleep(Duration.millis(options.errorBackoffMs ?? config.consumerErrorBackoffMs)),
            ),
          ),
        ),
      ),
    step: () => undefined,
  });

export const makeEffectConsumerFromPubSub = Effect.fn("makeEffectConsumerFromPubSub")(function* <T, E, R>(
  pubsub: PubSubService,
  config: MessagingRuntimeConfig,
  options: EffectConsumerOptions<T, E, R>,
  flow: FlowContext<R>,
) {
  const createOptions: CreateConsumerOptions<T> = {
    topic: options.topic,
    subscription: options.subscription,
    ...(options.initialPosition === undefined ? {} : { initialPosition: options.initialPosition }),
    ...(options.schema === undefined ? {} : { schema: options.schema }),
  };
  const concurrency = Math.max(1, options.concurrency ?? 1);
  const workerIndexes = Array.from({ length: concurrency }, (_value, index) => index);
  const workerConfig = {
    ...config,
    rateLimitRetryMs: options.rateLimitRetryMs ?? config.rateLimitRetryMs,
    rateLimitTimeoutMs: options.rateLimitTimeoutMs ?? config.rateLimitTimeoutMs,
  };
  const workers = yield* Effect.forEach(workerIndexes, () =>
    Effect.gen(function* () {
      const backend = yield* pubsub.createConsumer<T>(createOptions);
      const fiber = yield* consumerLoop(backend, options, flow, workerConfig).pipe(Effect.forkScoped);
      return { backend, fiber };
    }),
  );
  const stopped = yield* Ref.make(false);

  const stop = Effect.fn(`Consumer.stop:${options.topic}`)(function* () {
    const alreadyStopped = yield* Ref.getAndSet(stopped, true);
    if (alreadyStopped) return;

    yield* Effect.forEach(workers, (worker) => Fiber.interrupt(worker.fiber), { discard: true });
    yield* Effect.forEach(
      workers,
      (worker) => closeConsumerBackend(worker.backend, options.topic, options.subscription),
      { discard: true },
    );
  });

  yield* Effect.addFinalizer(() =>
    stop().pipe(
      Effect.catch((error) =>
        Effect.logError("[Consumer] Failed to stop consumer", {
          error: error.message,
          resource: error.resource,
          operation: error.operation,
        }),
      ),
    ),
  );

  return {
    fibers: workers.map((worker) => worker.fiber),
    stop: stop(),
  } satisfies EffectConsumer;
});

const dispatchResponseLoop = <T>(
  backend: BackendConsumer<T>,
  responseTopic: string,
  responses: EffectPubSub.PubSub<ResponseEnvelope<T>>,
  config: MessagingRuntimeConfig,
): Effect.Effect<void> =>
  Effect.whileLoop({
    while: () => true,
    body: () =>
      receiveMessage(backend, responseTopic, config.consumerReceiveTimeoutMs).pipe(
        Effect.flatMap((message) => {
          if (message === null) {
            return Effect.sleep(Duration.millis(config.consumerReceiveTimeoutMs));
          }

          const id = message.properties().id;
          return Effect.gen(function* () {
            if (id !== undefined) {
              yield* EffectPubSub.publish(responses, {
                id,
                value: message.value(),
              });
            }
            yield* acknowledgeMessage(backend, message, responseTopic);
          });
        }),
        Effect.catch((error) =>
          Effect.logError("[RequestResponse] Response dispatch failed", {
            error: error.message,
            topic: responseTopic,
          }).pipe(Effect.flatMap(() => Effect.sleep(Duration.millis(config.consumerErrorBackoffMs)))),
        ),
      ),
    step: () => undefined,
  });

const waitForResponse = Effect.fn("waitForResponse")(function* <TRes, E, R>(
  subscription: EffectPubSub.Subscription<ResponseEnvelope<TRes>>,
  id: string,
  options: EffectRequestOptions<TRes, E, R> | undefined,
) {
  const response = yield* Stream.fromSubscription(subscription).pipe(
    Stream.filterMapEffect((candidate) => {
      if (candidate.id !== id) {
        return Effect.succeed(Result.fail(undefined));
      }

      if (options?.recipient === undefined) {
        return Effect.succeed(Result.succeed(candidate.value));
      }

      return options.recipient(candidate.value).pipe(
        Effect.map((complete) =>
          complete
            ? Result.succeed(candidate.value)
            : Result.fail(undefined)
        ),
      );
    }),
    Stream.runHead,
  );

  return yield* O.match(response, {
    onNone: () => Effect.never,
    onSome: Effect.succeed,
  });
});

export const makeEffectRequestResponseFromPubSub = Effect.fn("makeEffectRequestResponseFromPubSub")(function* <
  TReq,
  TRes,
>(
  pubsub: PubSubService,
  config: MessagingRuntimeConfig,
  options: EffectRequestResponseOptions<TReq, TRes>,
) {
  const producerOptions: CreateProducerOptions<TReq> = options.requestSchema === undefined
    ? { topic: options.requestTopic }
    : { topic: options.requestTopic, schema: options.requestSchema };
  const producerBackend = yield* pubsub.createProducer<TReq>(producerOptions);
  const producer = makeEffectProducerHandle<TReq>(producerBackend, {
    topic: options.requestTopic,
    ...(options.requestSchema === undefined ? {} : { schema: options.requestSchema }),
  });
  const createOptions: CreateConsumerOptions<TRes> = {
    topic: options.responseTopic,
    subscription: options.subscription,
    ...(options.responseSchema === undefined ? {} : { schema: options.responseSchema }),
  };
  const backend = yield* pubsub.createConsumer<TRes>(createOptions);
  const responses = yield* EffectPubSub.unbounded<ResponseEnvelope<TRes>>();
  const stoppedSignal = yield* Deferred.make<never, MessagingLifecycleError>();
  const fiber = yield* dispatchResponseLoop(backend, options.responseTopic, responses, config).pipe(Effect.forkScoped);
  let stopped = false;

  const stop = Effect.fn(`RequestResponse.stop:${options.requestTopic}`)(function* () {
    if (stopped) return;
    stopped = true;
    yield* Deferred.fail(
      stoppedSignal,
      messagingLifecycleError(`${options.requestTopic}:${options.responseTopic}`, "stop", "RequestResponse stopped"),
    ).pipe(Effect.ignore);
    yield* EffectPubSub.shutdown(responses).pipe(Effect.ignore);
    yield* Fiber.interrupt(fiber);
    yield* producer.close;
    yield* closeConsumerBackend(backend, options.responseTopic, options.subscription);
  });

  yield* Effect.addFinalizer(() =>
    stop().pipe(
      Effect.catch((error) =>
        Effect.logError("[RequestResponse] Failed to stop runtime", {
          error: error.message,
        }),
      ),
    ),
  );

  return {
    request: <E = never, R = never>(
      request: TReq,
      requestOptions?: EffectRequestOptions<TRes, E, R>,
    ) => {
      const id = randomUUID();
      const timeoutMs = requestOptions?.timeoutMs ?? config.requestTimeoutMs;

      return Effect.scoped(
        Effect.gen(function* () {
          const subscription = yield* EffectPubSub.subscribe(responses);
          yield* producer.send(id, request);
          const result = yield* waitForResponse(subscription, id, requestOptions).pipe(
            Effect.raceFirst(Deferred.await(stoppedSignal)),
            Effect.timeoutOption(Duration.millis(timeoutMs)),
          );
          return yield* O.match(result, {
            onNone: () => Effect.fail(messagingTimeoutError("request-response", timeoutMs)),
            onSome: Effect.succeed,
          });
        }),
      );
    },
    stop: stop(),
  } satisfies EffectRequestResponse<TReq, TRes>;
});

export function makeProducerFactoryService(pubsub: PubSubService): ProducerFactoryService {
  return {
    make: Effect.fn("ProducerFactory.make")(<T>(options: EffectProducerOptions<T>) =>
      makeEffectProducerFromPubSub<T>(pubsub, options),
    ),
  };
}

export function makeConsumerFactoryService(
  pubsub: PubSubService,
  config: MessagingRuntimeConfig,
): ConsumerFactoryService {
  return {
    run: Effect.fn("ConsumerFactory.run")(<T, E = never, R = never>(
      options: EffectConsumerOptions<T, E, R>,
      flow: FlowContext<R>,
    ) =>
      makeEffectConsumerFromPubSub(pubsub, config, options, flow),
    ),
  };
}

export function makeRequestResponseFactoryService(
  pubsub: PubSubService,
  config: MessagingRuntimeConfig,
): RequestResponseFactoryService {
  return {
    make: Effect.fn("RequestResponseFactory.make")(<TReq, TRes>(
      options: EffectRequestResponseOptions<TReq, TRes>,
    ) => makeEffectRequestResponseFromPubSub<TReq, TRes>(pubsub, config, options)),
  };
}

export const ProducerFactoryLive = Layer.effect(
  ProducerFactory,
  Effect.gen(function* () {
    const pubsub = yield* PubSub;
    return ProducerFactory.of(makeProducerFactoryService(pubsub));
  }),
);

export const ConsumerFactoryLive = Layer.effect(
  ConsumerFactory,
  Effect.gen(function* () {
    const pubsub = yield* PubSub;
    const config = yield* loadMessagingRuntimeConfig();
    return ConsumerFactory.of(makeConsumerFactoryService(pubsub, config));
  }),
);

export const RequestResponseFactoryLive = Layer.effect(
  RequestResponseFactory,
  Effect.gen(function* () {
    const pubsub = yield* PubSub;
    const config = yield* loadMessagingRuntimeConfig();
    return RequestResponseFactory.of(makeRequestResponseFactoryService(pubsub, config));
  }),
);

export const runFlowRuntimeScoped = Effect.fn("FlowRuntime.run")(function* <Requirements = never>(
  flow: Flow<Requirements>,
) {
  yield* flow.startEffect().pipe(
    Effect.mapError((error) => flowRuntimeError(flow.name, "start", error)),
  );
  yield* Effect.addFinalizer(() =>
    Effect.sync(() => {
      flow.clearResources();
    }),
  );
});

export const FlowRuntimeLive = Layer.succeed(
  FlowRuntime,
  FlowRuntime.of({
    run: runFlowRuntimeScoped,
  }),
);

export const MessagingRuntimeLive = Layer.mergeAll(
  ProducerFactoryLive,
  ConsumerFactoryLive,
  RequestResponseFactoryLive,
  FlowRuntimeLive,
);

export const runEffectProducerScoped = Effect.fn("runEffectProducerScoped")(function* <T>(
  options: EffectProducerOptions<T>,
) {
  const pubsub = yield* PubSub;
  return yield* makeEffectProducerFromPubSub<T>(pubsub, options);
});

export const runEffectConsumerScoped = Effect.fn("runEffectConsumerScoped")(function* <T, E = never, R = never>(
  options: EffectConsumerOptions<T, E, R>,
  flow: FlowContext<R>,
) {
  const pubsub = yield* PubSub;
  const config = yield* loadMessagingRuntimeConfig();
  return yield* makeEffectConsumerFromPubSub(pubsub, config, options, flow);
});

export const runEffectRequestResponseScoped = Effect.fn("runEffectRequestResponseScoped")(function* <TReq, TRes>(
  options: EffectRequestResponseOptions<TReq, TRes>,
) {
  const pubsub = yield* PubSub;
  const config = yield* loadMessagingRuntimeConfig();
  return yield* makeEffectRequestResponseFromPubSub<TReq, TRes>(pubsub, config, options);
});

export const runFlowScoped = Effect.fn("runFlowScoped")(function* (
  flow: Flow,
) {
  yield* runFlowRuntimeScoped(flow);
});