fix(webapp,run-engine): stop batchTriggerAndWait hanging when item streaming never completes

.server-changes/batch-item-streaming-rate-limit.md

@@ -0,0 +1,6 @@
+---
+area: webapp
+type: fix
+---
+
+Batch triggers no longer fail to start their runs when an environment is under heavy API load. If a batch still can't finish being created, `batchTriggerAndWait` now fails with an error instead of leaving the parent run waiting forever, and the batches page says so rather than reporting that it resumed.

apps/webapp/app/env.server.ts

@@ -879,6 +879,20 @@ const EnvironmentSchema = z
     BATCH_RATE_LIMIT_MAX: z.coerce.number().int().default(1200),
     BATCH_RATE_LIMIT_REFILL_INTERVAL: z.string().default("10s"),
     BATCH_CONCURRENCY_LIMIT_DEFAULT: z.coerce.number().int().default(5),
+    /**
+     * How long a created batch may remain unsealed before the seal-timeout reaper
+     * aborts it and resumes any blocked parent with an error. Must exceed the SDK's
+     * worst-case stream-retry budget (maxAttempts x server request timeout).
+     * Doubles as the TTL of the phase 2 streaming grant, so the grant and the reaper
+     * always agree on how long a batch is allowed to be sealing.
+     */
+    BATCH_SEAL_TIMEOUT_MS: z.coerce.number().int().positive().default(1_800_000),
+    /**
+     * Number of phase 2 (`POST /api/v3/batches/:id/items`) requests a created batch is
+     * granted, exempt from the general API rate limit. Sized above the SDK's stream
+     * maxAttempts so a batch admitted by the batch limiter can always finish streaming.
+     */
+    BATCH_STREAM_GRANT_ATTEMPTS: z.coerce.number().int().positive().default(10),
 
     REALTIME_STREAM_VERSION: z.enum(["v1", "v2"]).default("v1"),
     REALTIME_STREAM_MAX_LENGTH: z.coerce.number().int().default(1000),

apps/webapp/app/routes/resources.batches.$batchId.check-completion.ts

@@ -42,6 +42,20 @@ export const action: ActionFunction = async ({ request, params }) => {
     return redirectWithErrorMessage(safeRedirectUrl, request, "Batch not found");
   }
 
+  const batch = await runStore.findBatchTaskRunById(ownedBatchRunId);
+
+  if (!batch) {
+    return redirectWithErrorMessage(safeRedirectUrl, request, "Batch not found");
+  }
+
+  if (!batch.sealed) {
+    return redirectWithErrorMessage(
+      safeRedirectUrl,
+      request,
+      "This batch was never finished being created, so it can't be resumed. Please get in touch and we'll recover it for you."
+    );
+  }
+
   try {
     // v3 (engine V1) is retired; finalize the batch through the v2 completion path (no-op if not ready).
     await tryCompleteBatchV3(ownedBatchRunId, prisma, true);

apps/webapp/app/runEngine/concerns/batchStreamGrants.server.ts

@@ -0,0 +1,105 @@
+import { createRedisClient, type RedisClient, type RedisWithClusterOptions } from "~/redis.server";
+import { logger } from "~/services/logger.server";
+
+export type BatchStreamGrantsOptions = {
+  redis: RedisWithClusterOptions;
+  /** How many phase 2 requests a created batch is allowed. */
+  attempts: number;
+  /** How long the grant survives, matching how long a batch may legitimately be sealing. */
+  ttlMs: number;
+};
+
+const KEY_PREFIX = "batch-stream-grant:";
+
+/**
+ * Admission for phase 2 of the 2-phase batch API.
+ *
+ * Phase 1 (`POST /api/v3/batches`) already passes its own batch rate limiter, which fixes
+ * the batch's `expectedCount` and blocks the parent run on the batch's waitpoint. Phase 2
+ * (`POST /api/v3/batches/:id/items`) is the only thing that can seal that batch, so having
+ * the general API limiter reject it strands the batch and the parent with it.
+ *
+ * Phase 1 therefore mints a bounded grant, and phase 2 spends it to bypass the general
+ * limiter. Admission stays a single decision made in phase 1, but the bypass is capped at
+ * `attempts` requests per batch rather than being unconditional.
+ */
+export class BatchStreamGrants {
+  private readonly redis: RedisClient;
+
+  constructor(private readonly options: BatchStreamGrantsOptions) {
+    this.redis = createRedisClient("batchStreamGrants", options.redis);
+    this.#registerCommands();
+  }
+
+  /**
+   * Grant a newly created batch its phase 2 budget. Never throws: a batch that fails to get
+   * a grant still works, it just falls back to the general rate limiter for streaming.
+   */
+  async mint(environmentId: string, batchId: string): Promise<void> {
+    try {
+      await this.redis.set(
+        this.#key(environmentId, batchId),
+        this.options.attempts,
+        "PX",
+        this.options.ttlMs
+      );
+    } catch (error) {
+      logger.warn("BatchStreamGrants: failed to mint grant", {
+        batchId,
+        error: error instanceof Error ? error.message : String(error),
+      });
+    }
+  }
+
+  /**
+   * Consume one phase 2 request from the batch's grant.
+   *
+   * Returns false when there is no grant, when the budget is spent, or when Redis is
+   * unreachable, so the caller falls back to the general rate limiter rather than opening
+   * an unbounded bypass.
+   */
+  async spend(environmentId: string, batchId: string): Promise<boolean> {
+    try {
+      // @ts-expect-error - Custom command defined via defineCommand
+      const remaining = (await this.redis.spendBatchStreamGrant(
+        this.#key(environmentId, batchId)
+      )) as number;
+
+      return remaining >= 0;
+    } catch (error) {
+      logger.warn("BatchStreamGrants: failed to spend grant", {
+        batchId,
+        error: error instanceof Error ? error.message : String(error),
+      });
+
+      return false;
+    }
+  }
+
+  async quit(): Promise<void> {
+    await this.redis.quit();
+  }
+
+  /**
+   * Scoped to the environment as well as the batch, so a caller authenticated against a
+   * different environment can never spend this batch's grant even if they know its id.
+   */
+  #key(environmentId: string, batchId: string): string {
+    return `${KEY_PREFIX}${environmentId}:${batchId}`;
+  }
+
+  #registerCommands(): void {
+    this.redis.defineCommand("spendBatchStreamGrant", {
+      numberOfKeys: 1,
+      lua: `
+local remaining = tonumber(redis.call('GET', KEYS[1]))
+
+if not remaining or remaining <= 0 then
+  return -1
+end
+
+return redis.call('DECR', KEYS[1])
+      `,
+    });
+  }
+}

apps/webapp/app/runEngine/concerns/batchStreamGrantsInstance.server.ts

@@ -0,0 +1,20 @@
+import { env } from "~/env.server";
+import { singleton } from "~/utils/singleton";
+import { BatchStreamGrants } from "./batchStreamGrants.server";
+
+export const batchStreamGrants = singleton(
+  "batchStreamGrants",
+  () =>
+    new BatchStreamGrants({
+      redis: {
+        port: env.RATE_LIMIT_REDIS_PORT,
+        host: env.RATE_LIMIT_REDIS_HOST,
+        username: env.RATE_LIMIT_REDIS_USERNAME,
+        password: env.RATE_LIMIT_REDIS_PASSWORD,
+        tlsDisabled: env.RATE_LIMIT_REDIS_TLS_DISABLED === "true",
+        clusterMode: env.RATE_LIMIT_REDIS_CLUSTER_MODE_ENABLED === "1",
+      },
+      attempts: env.BATCH_STREAM_GRANT_ATTEMPTS,
+      ttlMs: env.BATCH_SEAL_TIMEOUT_MS,
+    })
+);

apps/webapp/app/runEngine/services/createBatch.server.ts

@@ -4,6 +4,8 @@ import { RunId } from "@trigger.dev/core/v3/isomorphic";
 import { type BatchTaskRun, Prisma } from "@trigger.dev/database";
 import { Evt } from "evt";
 import { prisma, type PrismaClientOrTransaction } from "~/db.server";
+import { env } from "~/env.server";
+import { batchStreamGrants } from "../concerns/batchStreamGrantsInstance.server";
 import type { AuthenticatedEnvironment } from "~/services/apiAuth.server";
 import { logger } from "~/services/logger.server";
 import { controlPlaneResolver } from "~/v3/runOpsMigration/controlPlaneResolver.server";
@@ -120,8 +122,15 @@ export class CreateBatchService extends WithRunEngine {
 
           this.onBatchTaskRunCreated.post(batch);
 
+          await batchStreamGrants.mint(environment.id, friendlyId);
+
           // Block parent run if this is a batchTriggerAndWait
           if (body.parentRunId && body.resumeParentOnCompletion) {
+            await this._engine.scheduleExpireBatch({
+              batchId: batch.id,
+              availableAt: new Date(Date.now() + env.BATCH_SEAL_TIMEOUT_MS),
+            });
+
             await this._engine.blockRunWithCreatedBatch({
               runId: RunId.fromFriendlyId(body.parentRunId),
               batchId: batch.id,

apps/webapp/app/services/apiRateLimit.server.ts

@@ -1,8 +1,12 @@
+import { tryCatch } from "@trigger.dev/core/v3";
 import { env } from "~/env.server";
+import { batchStreamGrants } from "~/runEngine/concerns/batchStreamGrantsInstance.server";
 import { authenticateAuthorizationHeader } from "./apiAuth.server";
 import { authorizationRateLimitMiddleware } from "./authorizationRateLimitMiddleware.server";
 import type { Duration } from "./rateLimiter.server";
 
+const BATCH_STREAM_ITEMS_PATH = /^\/api\/v3\/batches\/([^/]+)\/items$/;
+
 export const apiRateLimiter = authorizationRateLimitMiddleware({
   redis: {
     port: env.RATE_LIMIT_REDIS_PORT,
@@ -75,6 +79,32 @@ export const apiRateLimiter = authorizationRateLimitMiddleware({
     /^\/api\/v2\/packets\//,
     /^\/api\/v1\/sessions\/[^/]+\/snapshot-url$/,
   ],
+  bypass: async (req) => {
+    const match = BATCH_STREAM_ITEMS_PATH.exec(req.path);
+
+    if (!match) {
+      return false;
+    }
+
+    const batchFriendlyId = match[1];
+    const authorizationValue = req.headers.authorization;
+
+    if (!batchFriendlyId || !authorizationValue) {
+      return false;
+    }
+
+    const [authError, authenticated] = await tryCatch(
+      authenticateAuthorizationHeader(authorizationValue, {
+        allowPublicKey: true,
+      })
+    );
+
+    if (authError || !authenticated || !authenticated.ok) {
+      return false;
+    }
+
+    return batchStreamGrants.spend(authenticated.environment.id, batchFriendlyId);
+  },
   log: {
     rejections: env.API_RATE_LIMIT_REJECTION_LOGS_ENABLED === "1",
     requests: env.API_RATE_LIMIT_REQUEST_LOGS_ENABLED === "1",

apps/webapp/app/services/authorizationRateLimitMiddleware.server.ts

@@ -5,7 +5,6 @@ import { Ratelimit } from "@upstash/ratelimit";
 import type { Request as ExpressRequest, Response as ExpressResponse, NextFunction } from "express";
 import { createHash } from "node:crypto";
 import { z } from "zod";
-import { env } from "~/env.server";
 import type { RedisWithClusterOptions } from "~/redis.server";
 import { logger } from "./logger.server";
 import type { Duration, Limiter } from "./rateLimiter.server";
@@ -56,10 +55,17 @@ export type RateLimiterConfig = z.infer<typeof RateLimiterConfig>;
 type LimitConfigOverrideFunction = (authorizationValue: string) => Promise<unknown>;
 
 type Options = {
-  redis?: RedisWithClusterOptions;
+  redis: RedisWithClusterOptions;
   keyPrefix: string;
   pathMatchers: (RegExp | string)[];
   pathWhiteList?: (RegExp | string)[];
+  /**
+   * Escape hatch for requests that can only be admitted by consulting state, rather than by
+   * matching a path. Runs after the authorization header check, so an unauthenticated
+   * request is still rejected, and only skips the rate limit itself. Must not throw: a
+   * bypass that cannot decide should return false and let the limiter apply.
+   */
+  bypass?: (req: ExpressRequest) => Promise<boolean>;
   defaultLimiter: RateLimiterConfig;
   limiterConfigOverride?: LimitConfigOverrideFunction;
   limiterCache?: {
@@ -151,6 +157,7 @@ export function authorizationRateLimitMiddleware({
   defaultLimiter,
   pathMatchers,
   pathWhiteList = [],
+  bypass,
   log = {
     rejections: true,
     requests: true,
@@ -176,16 +183,7 @@ export function authorizationRateLimitMiddleware({
     }),
   });
 
-  const redisClient = createRedisRateLimitClient(
-    redis ?? {
-      port: env.RATE_LIMIT_REDIS_PORT,
-      host: env.RATE_LIMIT_REDIS_HOST,
-      username: env.RATE_LIMIT_REDIS_USERNAME,
-      password: env.RATE_LIMIT_REDIS_PASSWORD,
-      tlsDisabled: env.RATE_LIMIT_REDIS_TLS_DISABLED === "true",
-      clusterMode: env.RATE_LIMIT_REDIS_CLUSTER_MODE_ENABLED === "1",
-    }
-  );
+  const redisClient = createRedisRateLimitClient(redis);
 
   return async (req: ExpressRequest, res: ExpressResponse, next: NextFunction) => {
     if (log.requests) {
@@ -247,6 +245,26 @@ export function authorizationRateLimitMiddleware({
       );
     }
 
+    if (bypass) {
+      let bypassed = false;
+
+      try {
+        bypassed = await bypass(req);
+      } catch (error) {
+        logger.warn(`RateLimiter (${keyPrefix}): bypass threw, applying the limit`, {
+          path: req.path,
+          error: error instanceof Error ? error.message : String(error),
+        });
+      }
+
+      if (bypassed) {
+        if (log.requests) {
+          logger.info(`RateLimiter (${keyPrefix}): bypassed ${req.path}`);
+        }
+        return next();
+      }
+    }
+
     const hash = createHash("sha256");
     hash.update(authorizationValue);
     const hashedAuthorizationValue = hash.digest("hex");

apps/webapp/app/v3/runOpsMigration/unblockRouteCatalog.ts

@@ -65,6 +65,12 @@ export const UNBLOCK_ROUTES: readonly UnblockRoute[] = [
     site: BATCH_SYSTEM,
     symbol: "#tryCompleteBatch",
   },
+  {
+    id: "batch.expireBatch",
+    kind: "RUN",
+    site: BATCH_SYSTEM,
+    symbol: "expireBatch",
+  },
   {
     id: "ttl.expireRun",
     kind: "RUN",