feat(webapp,redis-worker): mutateWithFallback helper (Phase B5)

d-cs · d-cs · commit dea1c7c0d99a · 2026-05-20T15:51:05.000+01:00
Composes PG-first (replica) lookup, MollifierBuffer.mutateSnapshot,
and writer-side spin-wait into the Q3 wait-and-bounce flow. Returns
a discriminated outcome rather than throwing Response, so the helper
stays route-agnostic and unit-testable. Phase C mutation endpoints
(tags, metadata-put, reschedule, cancel) consume this in upcoming
commits.

Wait knobs default to safetyNetMs=2000, pollStepMs=20, pgTimeoutMs=50
per Q3. Each PG poll is bounded by pgTimeoutMs via Promise.race so
a slow query can't burn the whole safety-net budget. Abort signal is
respected between polls (callers should pass getRequestAbortSignal()
when running in a request handler).

Also exports SnapshotPatch and MutateSnapshotResult from
@trigger.dev/redis-worker so webapp consumers can type-check their
callers of mutateSnapshot.
diff --git a/.changeset/mollifier-buffer-export-mutate-types.md b/.changeset/mollifier-buffer-export-mutate-types.md
@@ -0,0 +1,5 @@
+---
+"@trigger.dev/redis-worker": patch
+---
+
+Export `SnapshotPatch` and `MutateSnapshotResult` types from `@trigger.dev/redis-worker` so webapp consumers can type-check their callers of `MollifierBuffer.mutateSnapshot`.
diff --git a/.server-changes/mollifier-mutate-with-fallback-helper.md b/.server-changes/mollifier-mutate-with-fallback-helper.md
@@ -0,0 +1,6 @@
+---
+area: webapp
+type: improvement
+---
+
+Add `mutateWithFallback` helper in `app/v3/mollifier/mutateWithFallback.server.ts`. Composes PG-first (replica) lookup, `MollifierBuffer.mutateSnapshot`, and writer-side spin-wait into the Q3 wait-and-bounce flow. Returns a discriminated outcome (`pg` / `snapshot` / `not_found` / `timed_out`) without throwing Response objects, keeping the helper route-agnostic and unit-testable. Wait knobs (`safetyNetMs=2000`, `pollStepMs=20`, `pgTimeoutMs=50`) are overridable for tests. Each PG poll is bounded by `pgTimeoutMs` via `Promise.race` so a slow query can't burn the safety net. Phase C mutation endpoints (tags, metadata-put, reschedule, cancel) will consume this helper.
diff --git a/apps/webapp/app/v3/mollifier/mutateWithFallback.server.ts b/apps/webapp/app/v3/mollifier/mutateWithFallback.server.ts
@@ -0,0 +1,179 @@
+import type {
+  MollifierBuffer,
+  MutateSnapshotResult,
+  SnapshotPatch,
+} from "@trigger.dev/redis-worker";
+import type { TaskRun } from "@trigger.dev/database";
+import { prisma, $replica } from "~/db.server";
+import { logger } from "~/services/logger.server";
+import { getMollifierBuffer } from "./mollifierBuffer.server";
+
+// Wait/retry knobs per Q3 design. Exported for tests.
+export const DEFAULT_SAFETY_NET_MS = 2_000;
+export const DEFAULT_POLL_STEP_MS = 20;
+export const DEFAULT_PG_TIMEOUT_MS = 50;
+
+export type MutateWithFallbackInput<TResponse> = {
+  runId: string;
+  environmentId: string;
+  organizationId: string;
+  bufferPatch: SnapshotPatch;
+  // Called when a PG row exists (either replica-hit or post-wait writer-hit).
+  // Receives the full TaskRun shape and returns the customer-visible body.
+  pgMutation: (pgRow: TaskRun) => Promise<TResponse>;
+  // Called when the patch landed cleanly on the buffer snapshot. The
+  // drainer will see the patched payload on its next pop.
+  synthesisedResponse: () => TResponse;
+  abortSignal?: AbortSignal;
+  // Override defaults for tests.
+  safetyNetMs?: number;
+  pollStepMs?: number;
+  pgTimeoutMs?: number;
+  // Test injection.
+  getBuffer?: () => MollifierBuffer | null;
+  prismaWriter?: TaskRunReader;
+  prismaReplica?: TaskRunReader;
+  sleep?: (ms: number) => Promise<void>;
+  now?: () => number;
+};
+
+export type MutateWithFallbackOutcome<TResponse> =
+  | { kind: "pg"; response: TResponse }
+  | { kind: "snapshot"; response: TResponse }
+  | { kind: "not_found" }
+  | { kind: "timed_out" };
+
+// PG-first → buffer mutateSnapshot → wait-and-bounce. Implements the Q3
+// design (`_plans/2026-05-19-mollifier-mutation-race-design.md`). The
+// caller decides how to translate the outcome into an HTTP response —
+// this helper never throws Response objects so it remains route-agnostic
+// and unit-testable in isolation.
+export async function mutateWithFallback<TResponse>(
+  input: MutateWithFallbackInput<TResponse>,
+): Promise<MutateWithFallbackOutcome<TResponse>> {
+  const replica = input.prismaReplica ?? $replica;
+  const writer = input.prismaWriter ?? prisma;
+  const buffer = (input.getBuffer ?? getMollifierBuffer)();
+  const sleep = input.sleep ?? defaultSleep;
+  const now = input.now ?? Date.now;
+
+  // Path 1 — PG is already canonical.
+  const replicaRow = await findRunInPg(replica, input.runId, input.environmentId);
+  if (replicaRow) {
+    const response = await input.pgMutation(replicaRow);
+    return { kind: "pg", response };
+  }
+
+  if (!buffer) {
+    // No buffer configured (mollifier disabled or boot-time error). PG
+    // missed; nothing else to consult.
+    return { kind: "not_found" };
+  }
+
+  // Path 2 — buffer snapshot mutation.
+  const result: MutateSnapshotResult = await buffer.mutateSnapshot(
+    input.runId,
+    input.bufferPatch,
+  );
+
+  if (result === "applied_to_snapshot") {
+    return { kind: "snapshot", response: input.synthesisedResponse() };
+  }
+
+  if (result === "not_found") {
+    // Disambiguate a genuine 404 from a replica-lag miss: ask the writer
+    // directly. If the row just appeared post-drain we route through the
+    // PG mutation path.
+    const writerRow = await findRunInPg(writer, input.runId, input.environmentId);
+    if (writerRow) {
+      const response = await input.pgMutation(writerRow);
+      return { kind: "pg", response };
+    }
+    return { kind: "not_found" };
+  }
+
+  // result === "busy" — entry is DRAINING / FAILED / materialised. Wait
+  // for the drainer to terminate the entry into PG (success or
+  // SYSTEM_FAILURE) and route through pgMutation.
+  const safetyNetMs = input.safetyNetMs ?? DEFAULT_SAFETY_NET_MS;
+  const pollStepMs = input.pollStepMs ?? DEFAULT_POLL_STEP_MS;
+  const pgTimeoutMs = input.pgTimeoutMs ?? DEFAULT_PG_TIMEOUT_MS;
+  const deadline = now() + safetyNetMs;
+
+  while (now() < deadline) {
+    if (input.abortSignal?.aborted) {
+      return { kind: "timed_out" };
+    }
+
+    const row = await findRunInPgWithTimeout(
+      writer,
+      input.runId,
+      input.environmentId,
+      pgTimeoutMs,
+    );
+    if (row) {
+      const response = await input.pgMutation(row);
+      return { kind: "pg", response };
+    }
+
+    if (now() >= deadline) break;
+    await sleep(pollStepMs);
+  }
+
+  logger.warn("mollifier mutate-with-fallback: drainer resolution timed out", {
+    runId: input.runId,
+    safetyNetMs,
+  });
+  return { kind: "timed_out" };
+}
+
+// Structural reader interface — accepts both the writer (`prisma`) and the
+// replica (`$replica`), which differ slightly in their generated Prisma
+// types but share the findFirst surface used here.
+type TaskRunReader = {
+  taskRun: {
+    findFirst(args: {
+      where: { friendlyId: string; runtimeEnvironmentId: string };
+    }): Promise<TaskRun | null>;
+  };
+};
+
+async function findRunInPg(
+  client: TaskRunReader,
+  friendlyId: string,
+  environmentId: string,
+): Promise<TaskRun | null> {
+  return client.taskRun.findFirst({
+    where: { friendlyId, runtimeEnvironmentId: environmentId },
+  });
+}
+
+async function findRunInPgWithTimeout(
+  client: TaskRunReader,
+  friendlyId: string,
+  environmentId: string,
+  timeoutMs: number,
+): Promise<TaskRun | null> {
+  // One slow PG query shouldn't burn the whole safety-net budget.
+  // Promise.race against a timer; on timeout we treat the poll as a miss
+  // and the outer loop tries again on the next tick.
+  const timeoutToken = Symbol("pg-timeout");
+  let timeoutHandle: ReturnType<typeof setTimeout> | undefined;
+  const timeoutPromise = new Promise<typeof timeoutToken>((resolve) => {
+    timeoutHandle = setTimeout(() => resolve(timeoutToken), timeoutMs);
+  });
+  try {
+    const winner = await Promise.race([
+      findRunInPg(client, friendlyId, environmentId),
+      timeoutPromise,
+    ]);
+    if (winner === timeoutToken) return null;
+    return winner;
+  } finally {
+    if (timeoutHandle) clearTimeout(timeoutHandle);
+  }
+}
+
+function defaultSleep(ms: number): Promise<void> {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}
diff --git a/apps/webapp/test/mollifierMutateWithFallback.test.ts b/apps/webapp/test/mollifierMutateWithFallback.test.ts
@@ -0,0 +1,188 @@
+import { describe, expect, it, vi } from "vitest";
+
+vi.mock("~/db.server", () => ({
+  prisma: { taskRun: { findFirst: vi.fn(async () => null) } },
+  $replica: { taskRun: { findFirst: vi.fn(async () => null) } },
+}));
+
+import { mutateWithFallback } from "~/v3/mollifier/mutateWithFallback.server";
+import type { MollifierBuffer, MutateSnapshotResult } from "@trigger.dev/redis-worker";
+import type { TaskRun } from "@trigger.dev/database";
+
+type FindFirst = ReturnType<typeof vi.fn>;
+type PrismaStub = { taskRun: { findFirst: FindFirst } };
+
+function fakePrisma(rows: Array<TaskRun | null>): PrismaStub {
+  const fn = vi.fn();
+  for (const r of rows) fn.mockResolvedValueOnce(r);
+  fn.mockResolvedValue(null);
+  return { taskRun: { findFirst: fn } };
+}
+
+function bufferReturning(result: MutateSnapshotResult): MollifierBuffer {
+  return {
+    mutateSnapshot: vi.fn(async () => result),
+  } as unknown as MollifierBuffer;
+}
+
+const fakeRun = (overrides: Partial<TaskRun> = {}): TaskRun =>
+  ({
+    id: "pg_id",
+    friendlyId: "run_1",
+    runtimeEnvironmentId: "env_a",
+    ...overrides,
+  }) as TaskRun;
+
+const baseInput = {
+  runId: "run_1",
+  environmentId: "env_a",
+  organizationId: "org_1",
+  bufferPatch: { type: "append_tags" as const, tags: ["x"] },
+};
+
+describe("mutateWithFallback", () => {
+  it("hits replica → calls pgMutation, returns pg outcome", async () => {
+    const row = fakeRun();
+    const pgMutation = vi.fn(async () => "pg-response");
+    const synthesisedResponse = vi.fn(() => "snapshot-response");
+
+    const result = await mutateWithFallback({
+      ...baseInput,
+      pgMutation,
+      synthesisedResponse,
+      prismaReplica: fakePrisma([row]) as unknown as typeof import("~/db.server").$replica,
+      prismaWriter: fakePrisma([]) as unknown as typeof import("~/db.server").prisma,
+      getBuffer: () => bufferReturning("applied_to_snapshot"),
+    });
+
+    expect(result).toEqual({ kind: "pg", response: "pg-response" });
+    expect(pgMutation).toHaveBeenCalledWith(row);
+    expect(synthesisedResponse).not.toHaveBeenCalled();
+  });
+
+  it("replica miss + buffer applied_to_snapshot → synthesisedResponse", async () => {
+    const pgMutation = vi.fn(async () => "pg");
+    const result = await mutateWithFallback({
+      ...baseInput,
+      pgMutation,
+      synthesisedResponse: () => "snap",
+      prismaReplica: fakePrisma([null]) as unknown as typeof import("~/db.server").$replica,
+      prismaWriter: fakePrisma([]) as unknown as typeof import("~/db.server").prisma,
+      getBuffer: () => bufferReturning("applied_to_snapshot"),
+    });
+    expect(result).toEqual({ kind: "snapshot", response: "snap" });
+    expect(pgMutation).not.toHaveBeenCalled();
+  });
+
+  it("replica miss + buffer not_found + writer miss → not_found", async () => {
+    const result = await mutateWithFallback({
+      ...baseInput,
+      pgMutation: async () => "pg",
+      synthesisedResponse: () => "snap",
+      prismaReplica: fakePrisma([null]) as unknown as typeof import("~/db.server").$replica,
+      prismaWriter: fakePrisma([null]) as unknown as typeof import("~/db.server").prisma,
+      getBuffer: () => bufferReturning("not_found"),
+    });
+    expect(result).toEqual({ kind: "not_found" });
+  });
+
+  it("replica miss + buffer not_found + writer hit → pgMutation (replica-lag recovery)", async () => {
+    const row = fakeRun({ friendlyId: "run_1" });
+    const pgMutation = vi.fn(async () => "pg-recovered");
+    const result = await mutateWithFallback({
+      ...baseInput,
+      pgMutation,
+      synthesisedResponse: () => "snap",
+      prismaReplica: fakePrisma([null]) as unknown as typeof import("~/db.server").$replica,
+      prismaWriter: fakePrisma([row]) as unknown as typeof import("~/db.server").prisma,
+      getBuffer: () => bufferReturning("not_found"),
+    });
+    expect(result).toEqual({ kind: "pg", response: "pg-recovered" });
+    expect(pgMutation).toHaveBeenCalledWith(row);
+  });
+
+  it("replica miss + buffer busy + writer resolves mid-wait → pgMutation", async () => {
+    const row = fakeRun();
+    const pgMutation = vi.fn(async () => "pg-after-wait");
+    // Replica misses; writer misses twice, then hits.
+    const writer = fakePrisma([null, null, row]);
+    let nowValue = 0;
+    const result = await mutateWithFallback({
+      ...baseInput,
+      pgMutation,
+      synthesisedResponse: () => "snap",
+      prismaReplica: fakePrisma([null]) as unknown as typeof import("~/db.server").$replica,
+      prismaWriter: writer as unknown as typeof import("~/db.server").prisma,
+      getBuffer: () => bufferReturning("busy"),
+      sleep: async () => {
+        nowValue += 20;
+      },
+      now: () => nowValue,
+      safetyNetMs: 2000,
+      pollStepMs: 20,
+      pgTimeoutMs: 50,
+    });
+    expect(result).toEqual({ kind: "pg", response: "pg-after-wait" });
+    expect(pgMutation).toHaveBeenCalledWith(row);
+    // Writer should have been polled 3 times before the hit.
+    expect(writer.taskRun.findFirst).toHaveBeenCalledTimes(3);
+  });
+
+  it("replica miss + buffer busy + drainer never resolves → timed_out", async () => {
+    let nowValue = 0;
+    const result = await mutateWithFallback({
+      ...baseInput,
+      pgMutation: async () => "pg",
+      synthesisedResponse: () => "snap",
+      prismaReplica: fakePrisma([null]) as unknown as typeof import("~/db.server").$replica,
+      prismaWriter: fakePrisma([null, null, null, null, null]) as unknown as typeof import("~/db.server").prisma,
+      getBuffer: () => bufferReturning("busy"),
+      sleep: async () => {
+        nowValue += 20;
+      },
+      now: () => nowValue,
+      safetyNetMs: 60,
+      pollStepMs: 20,
+      pgTimeoutMs: 5,
+    });
+    expect(result).toEqual({ kind: "timed_out" });
+  });
+
+  it("abort signal during wait → timed_out without further polls", async () => {
+    const writer = fakePrisma([null, null, null]);
+    const controller = new AbortController();
+    let nowValue = 0;
+    const result = await mutateWithFallback({
+      ...baseInput,
+      pgMutation: async () => "pg",
+      synthesisedResponse: () => "snap",
+      prismaReplica: fakePrisma([null]) as unknown as typeof import("~/db.server").$replica,
+      prismaWriter: writer as unknown as typeof import("~/db.server").prisma,
+      getBuffer: () => bufferReturning("busy"),
+      sleep: async () => {
+        nowValue += 20;
+        controller.abort();
+      },
+      now: () => nowValue,
+      safetyNetMs: 2000,
+      pollStepMs: 20,
+      pgTimeoutMs: 5,
+      abortSignal: controller.signal,
+    });
+    expect(result).toEqual({ kind: "timed_out" });
+    // One poll happened before the sleep+abort.
+    expect(writer.taskRun.findFirst).toHaveBeenCalledTimes(1);
+  });
+
+  it("buffer is null (mollifier disabled) → not_found after replica miss", async () => {
+    const result = await mutateWithFallback({
+      ...baseInput,
+      pgMutation: async () => "pg",
+      synthesisedResponse: () => "snap",
+      prismaReplica: fakePrisma([null]) as unknown as typeof import("~/db.server").$replica,
+      prismaWriter: fakePrisma([]) as unknown as typeof import("~/db.server").prisma,
+      getBuffer: () => null,
+    });
+    expect(result).toEqual({ kind: "not_found" });
+  });
+});
diff --git a/packages/redis-worker/src/mollifier/index.ts b/packages/redis-worker/src/mollifier/index.ts
@@ -1,4 +1,9 @@
-export { MollifierBuffer, type MollifierBufferOptions } from "./buffer.js";
+export {
+  MollifierBuffer,
+  type MollifierBufferOptions,
+  type SnapshotPatch,
+  type MutateSnapshotResult,
+} from "./buffer.js";
 export {
   MollifierDrainer,
   type MollifierDrainerOptions,

-Original file line number
+Diff line change
@@ @@ -0,0 +1,5 @@ @@
 +---
 +"@trigger.dev/redis-worker": patch
 +---
++
 +Export `SnapshotPatch` and `MutateSnapshotResult` types from `@trigger.dev/redis-worker` so webapp consumers can type-check their callers of `MollifierBuffer.mutateSnapshot`.