fix(webapp,run-engine): scope batch stream grants to the environment and guard batch completion

ericallam · ericallam · commit 9ba38d1214a9 · 2026-07-27T16:06:05.000+01:00
Grants are now keyed by environment as well as batch, and are only spent once the
caller has been authenticated, so knowing a batch id is not enough to drain
another environment grant.

Batch completion no longer overwrites a terminal batch. An in-flight completion
that read the batch before it was aborted could previously flip it back to
completed and complete the waitpoint a second time with a success payload.
diff --git a/apps/webapp/app/runEngine/concerns/batchStreamGrants.server.ts b/apps/webapp/app/runEngine/concerns/batchStreamGrants.server.ts
@@ -35,9 +35,14 @@ export class BatchStreamGrants {
    * 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(batchId: string): Promise<void> {
+  async mint(environmentId: string, batchId: string): Promise<void> {
     try {
-      await this.redis.set(this.#key(batchId), this.options.attempts, "PX", this.options.ttlMs);
+      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,
@@ -53,10 +58,12 @@ export class BatchStreamGrants {
    * unreachable, so the caller falls back to the general rate limiter rather than opening
    * an unbounded bypass.
    */
-  async spend(batchId: string): Promise<boolean> {
+  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(batchId))) as number;
+      const remaining = (await this.redis.spendBatchStreamGrant(
+        this.#key(environmentId, batchId)
+      )) as number;
 
       return remaining >= 0;
     } catch (error) {
@@ -73,8 +80,12 @@ export class BatchStreamGrants {
     await this.redis.quit();
   }
 
-  #key(batchId: string): string {
-    return `${KEY_PREFIX}${batchId}`;
+  /**
+   * 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 {
diff --git a/apps/webapp/app/runEngine/services/createBatch.server.ts b/apps/webapp/app/runEngine/services/createBatch.server.ts
@@ -122,7 +122,7 @@ export class CreateBatchService extends WithRunEngine {
 
           this.onBatchTaskRunCreated.post(batch);
 
-          await batchStreamGrants.mint(friendlyId);
+          await batchStreamGrants.mint(environment.id, friendlyId);
 
           // Block parent run if this is a batchTriggerAndWait
           if (body.parentRunId && body.resumeParentOnCompletion) {
diff --git a/apps/webapp/app/services/apiRateLimit.server.ts b/apps/webapp/app/services/apiRateLimit.server.ts
@@ -86,12 +86,22 @@ export const apiRateLimiter = authorizationRateLimitMiddleware({
     }
 
     const batchFriendlyId = match[1];
+    const authorizationValue = req.headers.authorization;
 
-    if (!batchFriendlyId) {
+    if (!batchFriendlyId || !authorizationValue) {
       return false;
     }
 
-    return batchStreamGrants.spend(batchFriendlyId);
+    const authenticated = await authenticateAuthorizationHeader(authorizationValue, {
+      allowPublicKey: true,
+      allowJWT: true,
+    });
+
+    if (!authenticated || !authenticated.ok) {
+      return false;
+    }
+
+    return batchStreamGrants.spend(authenticated.environment.id, batchFriendlyId);
   },
   log: {
     rejections: env.API_RATE_LIMIT_REJECTION_LOGS_ENABLED === "1",
diff --git a/apps/webapp/test/authorizationRateLimitMiddleware.test.ts b/apps/webapp/test/authorizationRateLimitMiddleware.test.ts
@@ -415,81 +415,4 @@ describe.skipIf(process.env.GITHUB_ACTIONS)("authorizationRateLimitMiddleware",
       expect(configOverrideCalls).toBe(3);
     });
   });
-
-  describe("bypass", () => {
-    const exhaustedLimiter = {
-      type: "tokenBucket",
-      refillRate: 1,
-      interval: "1m",
-      maxTokens: 1,
-    } as const;
-
-    redisTest("lets a bypassed request through an exhausted limit", async ({ redisOptions }) => {
-      const rateLimitMiddleware = authorizationRateLimitMiddleware({
-        redis: { ...redisOptions, tlsDisabled: true },
-        keyPrefix: "test",
-        defaultLimiter: exhaustedLimiter,
-        pathMatchers: [/^\/api/],
-        bypass: async (req) => req.path === "/api/granted",
-      });
-
-      app.use(rateLimitMiddleware);
-      app.get("/api/granted", (req, res) => res.status(200).json({ message: "Granted" }));
-      app.get("/api/limited", (req, res) => res.status(200).json({ message: "Limited" }));
-
-      await request(app).get("/api/limited").set("Authorization", "Bearer test-token");
-      const limited = await request(app)
-        .get("/api/limited")
-        .set("Authorization", "Bearer test-token");
-      expect(limited.status).toBe(429);
-
-      const granted = await request(app)
-        .get("/api/granted")
-        .set("Authorization", "Bearer test-token");
-
-      expect(granted.status).toBe(200);
-      expect(granted.body).toEqual({ message: "Granted" });
-    });
-
-    redisTest("falls back to the limiter when the bypass declines", async ({ redisOptions }) => {
-      const rateLimitMiddleware = authorizationRateLimitMiddleware({
-        redis: { ...redisOptions, tlsDisabled: true },
-        keyPrefix: "test",
-        defaultLimiter: exhaustedLimiter,
-        pathMatchers: [/^\/api/],
-        bypass: async () => false,
-      });
-
-      app.use(rateLimitMiddleware);
-      app.get("/api/test", (req, res) => res.status(200).json({ message: "Success" }));
-
-      await request(app).get("/api/test").set("Authorization", "Bearer declined");
-      const response = await request(app).get("/api/test").set("Authorization", "Bearer declined");
-
-      expect(response.status).toBe(429);
-    });
-
-    redisTest("does not let the bypass skip authentication", async ({ redisOptions }) => {
-      let bypassCalled = false;
-
-      const rateLimitMiddleware = authorizationRateLimitMiddleware({
-        redis: { ...redisOptions, tlsDisabled: true },
-        keyPrefix: "test",
-        defaultLimiter: exhaustedLimiter,
-        pathMatchers: [/^\/api/],
-        bypass: async () => {
-          bypassCalled = true;
-          return true;
-        },
-      });
-
-      app.use(rateLimitMiddleware);
-      app.get("/api/test", (req, res) => res.status(200).json({ message: "Success" }));
-
-      const response = await request(app).get("/api/test");
-
-      expect(response.status).toBe(401);
-      expect(bypassCalled).toBe(false);
-    });
-  });
 });
diff --git a/apps/webapp/test/authorizationRateLimitMiddlewareBypass.test.ts b/apps/webapp/test/authorizationRateLimitMiddlewareBypass.test.ts
@@ -0,0 +1,92 @@
+import { redisTest } from "@internal/testcontainers";
+import { beforeEach, describe, expect, vi } from "vitest";
+
+vi.setConfig({ testTimeout: 30_000 });
+
+import type { Express } from "express";
+import express from "express";
+import request from "supertest";
+import { authorizationRateLimitMiddleware } from "../app/services/authorizationRateLimitMiddleware.server.js";
+
+const exhaustedLimiter = {
+  type: "tokenBucket",
+  refillRate: 1,
+  interval: "1m",
+  maxTokens: 1,
+} as const;
+
+describe("authorizationRateLimitMiddleware bypass", () => {
+  let app: Express;
+
+  beforeEach(() => {
+    app = express();
+  });
+
+  redisTest("lets a bypassed request through an exhausted limit", async ({ redisOptions }) => {
+    const rateLimitMiddleware = authorizationRateLimitMiddleware({
+      redis: { ...redisOptions, tlsDisabled: true },
+      keyPrefix: "test-bypass-allowed",
+      defaultLimiter: exhaustedLimiter,
+      pathMatchers: [/^\/api/],
+      bypass: async (req) => req.path === "/api/granted",
+    });
+
+    app.use(rateLimitMiddleware);
+    app.get("/api/granted", (req, res) => res.status(200).json({ message: "Granted" }));
+    app.get("/api/limited", (req, res) => res.status(200).json({ message: "Limited" }));
+
+    await request(app).get("/api/limited").set("Authorization", "Bearer test-token");
+    const limited = await request(app)
+      .get("/api/limited")
+      .set("Authorization", "Bearer test-token");
+    expect(limited.status).toBe(429);
+
+    const granted = await request(app)
+      .get("/api/granted")
+      .set("Authorization", "Bearer test-token");
+
+    expect(granted.status).toBe(200);
+    expect(granted.body).toEqual({ message: "Granted" });
+  });
+
+  redisTest("falls back to the limiter when the bypass declines", async ({ redisOptions }) => {
+    const rateLimitMiddleware = authorizationRateLimitMiddleware({
+      redis: { ...redisOptions, tlsDisabled: true },
+      keyPrefix: "test-bypass-declined",
+      defaultLimiter: exhaustedLimiter,
+      pathMatchers: [/^\/api/],
+      bypass: async () => false,
+    });
+
+    app.use(rateLimitMiddleware);
+    app.get("/api/test", (req, res) => res.status(200).json({ message: "Success" }));
+
+    await request(app).get("/api/test").set("Authorization", "Bearer declined");
+    const response = await request(app).get("/api/test").set("Authorization", "Bearer declined");
+
+    expect(response.status).toBe(429);
+  });
+
+  redisTest("does not let the bypass skip authentication", async ({ redisOptions }) => {
+    let bypassCalled = false;
+
+    const rateLimitMiddleware = authorizationRateLimitMiddleware({
+      redis: { ...redisOptions, tlsDisabled: true },
+      keyPrefix: "test-bypass-unauthenticated",
+      defaultLimiter: exhaustedLimiter,
+      pathMatchers: [/^\/api/],
+      bypass: async () => {
+        bypassCalled = true;
+        return true;
+      },
+    });
+
+    app.use(rateLimitMiddleware);
+    app.get("/api/test", (req, res) => res.status(200).json({ message: "Success" }));
+
+    const response = await request(app).get("/api/test");
+
+    expect(response.status).toBe(401);
+    expect(bypassCalled).toBe(false);
+  });
+});
diff --git a/apps/webapp/test/batchStreamGrants.test.ts b/apps/webapp/test/batchStreamGrants.test.ts
@@ -14,7 +14,9 @@ vi.mock("../app/services/logger.server", () => ({
 
 import { BatchStreamGrants } from "../app/runEngine/concerns/batchStreamGrants.server.js";
 
-describe.skipIf(process.env.GITHUB_ACTIONS)("BatchStreamGrants", () => {
+describe("BatchStreamGrants", () => {
+  const ENV = "env_1";
+
   redisTest("spends exactly the granted number of attempts", async ({ redisOptions }) => {
     const grants = new BatchStreamGrants({
       redis: { ...redisOptions, tlsDisabled: true },
@@ -23,13 +25,13 @@ describe.skipIf(process.env.GITHUB_ACTIONS)("BatchStreamGrants", () => {
     });
 
     try {
-      await grants.mint("batch_spend");
+      await grants.mint(ENV, "batch_spend");
 
-      expect(await grants.spend("batch_spend")).toBe(true);
-      expect(await grants.spend("batch_spend")).toBe(true);
-      expect(await grants.spend("batch_spend")).toBe(true);
-      expect(await grants.spend("batch_spend")).toBe(false);
-      expect(await grants.spend("batch_spend")).toBe(false);
+      expect(await grants.spend(ENV, "batch_spend")).toBe(true);
+      expect(await grants.spend(ENV, "batch_spend")).toBe(true);
+      expect(await grants.spend(ENV, "batch_spend")).toBe(true);
+      expect(await grants.spend(ENV, "batch_spend")).toBe(false);
+      expect(await grants.spend(ENV, "batch_spend")).toBe(false);
     } finally {
       await grants.quit();
     }
@@ -43,7 +45,7 @@ describe.skipIf(process.env.GITHUB_ACTIONS)("BatchStreamGrants", () => {
     });
 
     try {
-      expect(await grants.spend("batch_never_minted")).toBe(false);
+      expect(await grants.spend(ENV, "batch_never_minted")).toBe(false);
     } finally {
       await grants.quit();
     }
@@ -57,12 +59,12 @@ describe.skipIf(process.env.GITHUB_ACTIONS)("BatchStreamGrants", () => {
     });
 
     try {
-      await grants.mint("batch_a");
-      await grants.mint("batch_b");
+      await grants.mint(ENV, "batch_a");
+      await grants.mint(ENV, "batch_b");
 
-      expect(await grants.spend("batch_a")).toBe(true);
-      expect(await grants.spend("batch_a")).toBe(false);
-      expect(await grants.spend("batch_b")).toBe(true);
+      expect(await grants.spend(ENV, "batch_a")).toBe(true);
+      expect(await grants.spend(ENV, "batch_a")).toBe(false);
+      expect(await grants.spend(ENV, "batch_b")).toBe(true);
     } finally {
       await grants.quit();
     }
@@ -71,22 +73,39 @@ describe.skipIf(process.env.GITHUB_ACTIONS)("BatchStreamGrants", () => {
   redisTest("expires the grant so it cannot outlive the seal window", async ({ redisOptions }) => {
     const grants = new BatchStreamGrants({
       redis: { ...redisOptions, tlsDisabled: true },
-      attempts: 5,
+      attempts: 100_000,
       ttlMs: 150,
     });
 
     try {
-      await grants.mint("batch_expiring");
-      expect(await grants.spend("batch_expiring")).toBe(true);
+      await grants.mint(ENV, "batch_expiring");
+      expect(await grants.spend(ENV, "batch_expiring")).toBe(true);
 
       await vi.waitFor(
         async () => {
-          expect(await grants.spend("batch_expiring")).toBe(false);
+          expect(await grants.spend(ENV, "batch_expiring")).toBe(false);
         },
         { timeout: 5_000, interval: 50 }
       );
     } finally {
       await grants.quit();
     }
   });
+
+  redisTest("another environment cannot spend this batch's grant", async ({ redisOptions }) => {
+    const grants = new BatchStreamGrants({
+      redis: { ...redisOptions, tlsDisabled: true },
+      attempts: 1,
+      ttlMs: 60_000,
+    });
+
+    try {
+      await grants.mint(ENV, "batch_scoped");
+
+      expect(await grants.spend("env_intruder", "batch_scoped")).toBe(false);
+      expect(await grants.spend(ENV, "batch_scoped")).toBe(true);
+    } finally {
+      await grants.quit();
+    }
+  });
 });
diff --git a/internal-packages/run-engine/src/engine/systems/batchSystem.ts b/internal-packages/run-engine/src/engine/systems/batchSystem.ts
@@ -109,7 +109,7 @@ export class BatchSystem {
         return;
       }
 
-      const enqueuedCount = batch.successfulRunCount ?? 0;
+      const enqueuedCount = (batch.successfulRunCount ?? 0) + (batch.failedRunCount ?? 0);
       const error: TaskRunError = {
         type: "STRING_ERROR",
         raw:
@@ -190,21 +190,22 @@ export class BatchSystem {
 
       if (runs.every((r) => isFinalRunStatus(r.status))) {
         this.$.logger.debug("#tryCompleteBatch: All runs are completed", { batchId });
-        await this.$.runStore.updateBatchTaskRun(
+
+        const completed = await this.$.runStore.updateManyBatchTaskRun(
           {
-            where: {
-              id: batchId,
-            },
-            data: {
-              status: "COMPLETED",
-            },
-            select: {
-              id: true,
-            },
+            where: { id: batchId, status: { notIn: ["ABORTED", "COMPLETED"] } },
+            data: { status: "COMPLETED" },
           },
           this.$.prisma
         );
 
+        if (completed.count === 0) {
+          this.$.logger.debug("#tryCompleteBatch: batch already reached a terminal status", {
+            batchId,
+          });
+          return;
+        }
+
         //get waitpoint (if there is one)
         const waitpoint = await this.$.runStore.findWaitpoint(
           {
