triggerdotdev--trigger.dev/apps/webapp/app/services/realtime/runReader.server.ts

import {
  type Prisma,
  type PrismaClient,
  type PrismaClientOrTransaction,
  boundedIn,
} from "@trigger.dev/database";
import type { RunStore } from "@internal/run-store";
import { BoundedTtlCache } from "./boundedTtlCache";
import { RESERVED_COLUMNS, type RealtimeRunRow } from "./electricStreamProtocol.server";

/**
 * RunReader — the pluggable read half of the native-backend realtime feed: ClickHouse is filter-only
 * (resolves ids), Postgres always hydrates row columns. Owns the `RunHydrator` (by-id) and the
 * `RunListResolver` interface (the tag/list filter -> id-set seam, implemented over ClickHouse).
 */

/** The TaskRun columns the realtime feed projects (mirrors DEFAULT_ELECTRIC_COLUMNS). */
export const RUN_HYDRATOR_SELECT = {
  id: true,
  taskIdentifier: true,
  createdAt: true,
  updatedAt: true,
  startedAt: true,
  delayUntil: true,
  queuedAt: true,
  expiredAt: true,
  completedAt: true,
  friendlyId: true,
  number: true,
  isTest: true,
  status: true,
  usageDurationMs: true,
  costInCents: true,
  baseCostInCents: true,
  ttl: true,
  payload: true,
  payloadType: true,
  metadata: true,
  metadataType: true,
  output: true,
  outputType: true,
  runTags: true,
  error: true,
  realtimeStreams: true,
} satisfies Prisma.TaskRunSelect;

/** Columns hydrated regardless of `skipColumns`: `id` keys the row, `updatedAt` drives the offset and working-set diff. */
const ALWAYS_HYDRATED_COLUMNS = new Set<string>(["id", "updatedAt", ...RESERVED_COLUMNS]);

/** Project `RUN_HYDRATOR_SELECT` down to the columns the client didn't skip (plus
 * the always-needed ones). An empty skip set returns the full select unchanged. */
export function buildHydratorSelect(skipColumns: string[] = []): Prisma.TaskRunSelect {
  if (skipColumns.length === 0) {
    return RUN_HYDRATOR_SELECT;
  }
  const skip = new Set(skipColumns);
  const select: Record<string, boolean> = {};
  for (const column of Object.keys(RUN_HYDRATOR_SELECT)) {
    if (ALWAYS_HYDRATED_COLUMNS.has(column) || !skip.has(column)) {
      select[column] = true;
    }
  }
  return select as Prisma.TaskRunSelect;
}

export type RunListFilter = {
  organizationId: string;
  projectId: string;
  environmentId: string;
  /** Contains-ANY tag match (OR). Omit/empty for non-tag feeds. */
  tags?: string[];
  /** Restrict to a single batch (internal batch id) — the batch feed. */
  batchId?: string;
  /** Lower bound on createdAt (the tag-list feed pins this; batch omits it). */
  createdAtAfter?: Date;
  /** Hard cap on the result set so a broad filter can't unbound the snapshot. */
  limit: number;
};

/** Resolves a tag/list filter into the matching run id-set, filter-only (rows hydrated from Postgres by id afterward). ClickHouse impl in `clickHouseRunListResolver.server.ts`. */
export interface RunListResolver {
  resolveMatchingRunIds(filter: RunListFilter): Promise<string[]>;
}

export type RunHydratorOptions = {
  /** The Prisma client handed to the RunStore as the read client. Always Postgres. A branded
   * replica (`$replica`) keeps routed reads on each store's replica; an unbranded writer
   * (`prisma`) escalates them to each store's own primary. */
  readClient: Pick<PrismaClient, "taskRun">;
  /** RunStore the reads are routed through. */
  runStore: RunStore;
  /** Read-through cache TTL (ms) collapsing duplicate refetches for the same run. Set 0 to disable. Defaults to 250ms. */
  cacheTtlMs?: number;
  /** Hard cap on cache entries before expired entries are swept. */
  maxCacheEntries?: number;
};

const DEFAULT_CACHE_TTL_MS = 250;
const DEFAULT_MAX_CACHE_ENTRIES = 5_000;

/** Hydrates runs by id through the runStore seam (split routing lives in the store, below this file), projected to the realtime columns; concurrent same-run refetches are single-flighted + short-TTL cached. */
export class RunHydrator {
  readonly #inflight = new Map<string, Promise<RealtimeRunRow | null>>();
  readonly #cache: BoundedTtlCache<RealtimeRunRow | null>;
  readonly #cacheTtlMs: number;

  constructor(private readonly options: RunHydratorOptions) {
    this.#cacheTtlMs = options.cacheTtlMs ?? DEFAULT_CACHE_TTL_MS;
    this.#cache = new BoundedTtlCache(
      this.#cacheTtlMs,
      options.maxCacheEntries ?? DEFAULT_MAX_CACHE_ENTRIES
    );
  }

  async getRunById(environmentId: string, runId: string): Promise<RealtimeRunRow | null> {
    const key = `${environmentId}:${runId}`;

    if (this.#cacheTtlMs > 0) {
      // A cached null is a valid "run not found" hit; only undefined is a miss.
      const cached = this.#cache.get(key);
      if (cached !== undefined) {
        return cached;
      }
    }

    const existing = this.#inflight.get(key);
    if (existing) {
      return existing;
    }

    const promise = this.#fetch(environmentId, runId).finally(() => this.#inflight.delete(key));
    this.#inflight.set(key, promise);

    const row = await promise;

    if (this.#cacheTtlMs > 0) {
      this.#cache.set(key, row);
    }

    return row;
  }

  /** Hydrate many runs by id in one query (order not guaranteed); `skipColumns` projects the SELECT so dropped columns aren't shipped. */
  async hydrateByIds(
    environmentId: string,
    ids: string[],
    skipColumns: string[] = []
  ): Promise<RealtimeRunRow[]> {
    if (ids.length === 0) {
      return [];
    }
    const rows = await this.options.runStore.findRuns(
      {
        where: {
          runtimeEnvironmentId: environmentId,
          id: { in: boundedIn(ids) },
        },
        select: buildHydratorSelect(skipColumns),
      },
      this.options.readClient as PrismaClientOrTransaction
    );
    return rows as unknown as RealtimeRunRow[];
  }

  async #fetch(environmentId: string, runId: string): Promise<RealtimeRunRow | null> {
    const run = await this.options.runStore.findRun(
      {
        id: runId,
        runtimeEnvironmentId: environmentId,
      },
      { select: RUN_HYDRATOR_SELECT },
      this.options.readClient as PrismaClientOrTransaction
    );

    return (run ?? null) as RealtimeRunRow | null;
  }
}