From 61accf70642980653abe13a4a54a82df98687bc2 Mon Sep 17 00:00:00 2001 From: JSONbored <49853598+JSONbored@users.noreply.github.com> Date: Mon, 6 Jul 2026 06:06:12 -0700 Subject: [PATCH] fix(github): scope rate-limit admission checks to the caller's own bucket (#3809) shouldWaitForGitHubRateLimit read the globally-newest REST rate-limit observation across every installation and the shared public/registry token, regardless of which bucket the caller actually draws from. Since each installation and the shared public token have separate GitHub-side REST buckets, a healthy recent observation from one bucket could mask another bucket's real exhaustion, or an exhausted recent observation from one bucket could wrongly throttle a caller whose own bucket has headroom. Add an optional admissionKey parameter to shouldWaitForGitHubRateLimit and listLatestGitHubRateLimitObservations, and thread each caller's own already-computed admission key through at every call site where one is naturally available. One call site genuinely has no installation in scope yet (a pre-dispatch scheduler check) and is left unscoped, unchanged from before -- the real per-installation check still happens downstream once a specific job is dispatched. --- src/db/repositories.ts | 12 +++- src/github/backfill.ts | 15 +++-- src/github/rate-limit.ts | 12 +++- src/index.ts | 12 ++-- src/queue/dlq.ts | 9 ++- src/queue/processors.ts | 6 ++ src/upstream/commit.ts | 6 +- .../backfill-file-hydration-scoping.test.ts | 3 + test/unit/backfill.test.ts | 15 +++++ test/unit/index.test.ts | 3 +- test/unit/queue.test.ts | 24 ++++---- test/unit/rate-limit.test.ts | 57 ++++++++++++++++++- 12 files changed, 143 insertions(+), 31 deletions(-) diff --git a/src/db/repositories.ts b/src/db/repositories.ts index 2ecf070cb5..86c423ea62 100644 --- a/src/db/repositories.ts +++ b/src/db/repositories.ts @@ -1279,9 +1279,17 @@ export async function recordGitHubRateLimitObservation(env: Env, observation: Gi }); } -export async function listLatestGitHubRateLimitObservations(env: Env, limit = 50): Promise { +/** + * Latest observations, newest first. When `admissionKey` is given, scoped to ONLY that bucket (#audit-rate-scoping) + * — every managed installation and the separate shared public/registry token draw from DIFFERENT GitHub-side REST + * buckets, so an unscoped read can return the wrong bucket's row (e.g. a fresh public-token observation masking an + * exhausted installation bucket, or vice versa) purely because it happened to be the most recently written. + */ +export async function listLatestGitHubRateLimitObservations(env: Env, limit = 50, admissionKey?: string): Promise { const db = getDb(env.DB); - const rows = await db.select().from(githubRateLimitObservations).orderBy(desc(githubRateLimitObservations.observedAt)).limit(limit); + const rows = await (admissionKey !== undefined + ? db.select().from(githubRateLimitObservations).where(eq(githubRateLimitObservations.admissionKey, admissionKey)).orderBy(desc(githubRateLimitObservations.observedAt)).limit(limit) + : db.select().from(githubRateLimitObservations).orderBy(desc(githubRateLimitObservations.observedAt)).limit(limit)); return rows.map(toGitHubRateLimitObservationRecord); } diff --git a/src/github/backfill.ts b/src/github/backfill.ts index 3a42aa2f3f..b9114f7a41 100644 --- a/src/github/backfill.ts +++ b/src/github/backfill.ts @@ -71,7 +71,7 @@ import { shouldPublishReviewCheck, } from "../review/check-names"; import { buildReviewThreadBlocker, type ReviewThreadBlocker } from "../review/review-thread-findings"; -import { delayUntil, HISTORICAL_BACKFILL_RESERVED_HEADROOM, shouldWaitForGitHubRateLimit } from "./rate-limit"; +import { delayUntil, HISTORICAL_BACKFILL_RESERVED_HEADROOM, LOW_REST_RATE_LIMIT_REMAINING, shouldWaitForGitHubRateLimit } from "./rate-limit"; import { githubRateLimitAdmissionKeyForPublicToken, githubRateLimitAdmissionKeyForToken, @@ -507,7 +507,9 @@ export async function backfillRepositorySegment( const mode = options.mode ?? "light"; const token = await tokenForRepo(env, repo); const sourceKind: RepoSyncSegmentRecord["sourceKind"] = repo.installationId && token !== env.GITHUB_PUBLIC_TOKEN ? "installation" : "github"; - const resetAt = await shouldWaitForGitHubRateLimit(env); + // Scoped to the bucket this segment's OWN reads actually draw from (#audit-rate-scoping), not whichever bucket + // was most recently observed across every installation and the shared public token. + const resetAt = await shouldWaitForGitHubRateLimit(env, LOW_REST_RATE_LIMIT_REMAINING, repoAdmissionKeyForToken(env, repo, token)); if (resetAt) { const previous = await getRepoSyncSegment(env, repo.fullName, options.segment); const segment = await completeSegment(env, repo, options.segment, sourceKind, mode, nowIso(), { @@ -607,7 +609,9 @@ export async function backfillOpenPullRequestDetails( const mode = options.mode ?? "light"; const token = await tokenForRepo(env, repo); const sourceKind: RepoSyncSegmentRecord["sourceKind"] = repo.installationId && token !== env.GITHUB_PUBLIC_TOKEN ? "installation" : "github"; - const resetAt = await shouldWaitForGitHubRateLimit(env); + // Scoped to the bucket this segment's OWN reads actually draw from (#audit-rate-scoping), not whichever bucket + // was most recently observed across every installation and the shared public token. + const resetAt = await shouldWaitForGitHubRateLimit(env, LOW_REST_RATE_LIMIT_REMAINING, repoAdmissionKeyForToken(env, repo, token)); if (resetAt) { const previous = await getRepoSyncSegment(env, repo.fullName, "pull_request_files"); await env.JOBS.send({ type: "backfill-pr-details", requestedBy: "api", repoFullName: repo.fullName, ...repoInstallationPayload(repo), mode, cursor: options.cursor ?? 0 }, { delaySeconds: delayUntil(resetAt) }); @@ -1526,7 +1530,10 @@ async function hydrateMergedPullRequestFiles( .map((record) => record.number), ); const pending = merged.filter((pr) => !alreadyHydrated.has(pr.number)); - const resetAt = pending.length > 0 ? await shouldWaitForGitHubRateLimit(env, HISTORICAL_BACKFILL_RESERVED_HEADROOM) : undefined; + // #audit-rate-scoping: this function already receives (and forwards to its own GitHub reads below) the caller's + // admissionKey — thread it into the budget check too instead of checking whichever bucket was most recently + // observed across every installation and the shared public token. + const resetAt = pending.length > 0 ? await shouldWaitForGitHubRateLimit(env, HISTORICAL_BACKFILL_RESERVED_HEADROOM, admissionKey) : undefined; if (resetAt) warnings.push(`Historical merged PR file hydration deferred for ${pending.length} pull request(s): GitHub REST budget below the historical-backfill headroom floor (retry after ${resetAt}).`); const budgeted = resetAt ? new Set() : new Set(pending.slice(0, MERGED_PR_FILE_HYDRATION_BATCH_SIZE[mode]).map((pr) => pr.number)); await mapWithConcurrency(merged, concurrency, async (pr) => { diff --git a/src/github/rate-limit.ts b/src/github/rate-limit.ts index a74adb77f1..af9e5d3894 100644 --- a/src/github/rate-limit.ts +++ b/src/github/rate-limit.ts @@ -22,9 +22,15 @@ export const HISTORICAL_BACKFILL_RESERVED_HEADROOM = 300; /** The REST rate-limit reset time to wait until when the latest recorded REST budget is at/below `floor`, or * undefined when there is headroom, no usable observation, or the reset is already in the past. Reads the latest - * recorded observation (recordGitHubRateLimitObservation writes one per GitHub call) — no live GitHub call. */ -export async function shouldWaitForGitHubRateLimit(env: Env, floor: number = LOW_REST_RATE_LIMIT_REMAINING): Promise { - const observations = await listLatestGitHubRateLimitObservations(env, 10); + * recorded observation (recordGitHubRateLimitObservation writes one per GitHub call) — no live GitHub call. + * + * `admissionKey`, when given, scopes the read to that bucket ONLY (#audit-rate-scoping) — every installation and + * the shared public/registry token draw from DIFFERENT GitHub-side REST buckets, so checking the caller's own + * bucket instead of "whichever bucket was most recently observed" avoids one bucket's health masking or falsely + * throttling an unrelated one. Omitted (as at a few call sites with no natural single bucket to check, e.g. a + * pre-dispatch scheduler tick) preserves the prior globally-newest-observation behavior unchanged. */ +export async function shouldWaitForGitHubRateLimit(env: Env, floor: number = LOW_REST_RATE_LIMIT_REMAINING, admissionKey?: string): Promise { + const observations = await listLatestGitHubRateLimitObservations(env, 10, admissionKey); // Type-guard the find so `remaining` narrows to a number — null/undefined observations are excluded here, so the // headroom check below needs no further nullish guard. const rest = observations.find((observation): observation is typeof observation & { remaining: number } => observation.resource === "rest" && observation.remaining !== null && observation.remaining !== undefined); diff --git a/src/index.ts b/src/index.ts index af9930720e..cd9be9e4a8 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,6 +1,6 @@ import { createApp } from "./api/routes"; import { RateLimiter } from "./auth/rate-limit"; -import { delayUntil, shouldWaitForGitHubRateLimit, MAINTENANCE_RESERVED_HEADROOM } from "./github/rate-limit"; +import { delayUntil, shouldWaitForGitHubRateLimit, LOW_REST_RATE_LIMIT_REMAINING, MAINTENANCE_RESERVED_HEADROOM } from "./github/rate-limit"; import { processDlqBatch } from "./queue/dlq"; import { processJob } from "./queue/processors"; import { isOrbBrokerEnabled } from "./orb/broker"; @@ -8,6 +8,7 @@ import { isOpsEnabled } from "./review/ops-wire"; import { isRagEnabled } from "./review/rag-wire"; import { isSelfTuneEnabled } from "./review/selftune-wire"; import { + githubRateLimitAdmissionKeyForJob, isGitHubBudgetBackgroundJob, queueSnapshotBacklog, queueSnapshotFromBinding, @@ -54,7 +55,9 @@ export default { continue; } if (isGitHubBudgetBackgroundJob(message.body)) { - const resetAt = await shouldWaitForGitHubRateLimit(env, MAINTENANCE_RESERVED_HEADROOM).catch(() => undefined); + // Scoped to THIS job's own installation bucket (#audit-rate-scoping) — an unrelated installation's or + // the shared public token's budget must never defer (or wrongly clear) this job. + const resetAt = await shouldWaitForGitHubRateLimit(env, MAINTENANCE_RESERVED_HEADROOM, githubRateLimitAdmissionKeyForJob(message.body) ?? undefined).catch(() => undefined); if (resetAt) { console.log( JSON.stringify({ @@ -83,8 +86,9 @@ export default { ); // If the shared GitHub REST budget is exhausted, this failure is most likely a rate-limit — retry AFTER the // reset so a real webhook OUTLASTS a transient rate-limit window instead of burning its retries immediately - // and being dead-lettered (the surviving event-loss path). (#audit-rate-headroom) - const resetAt = await shouldWaitForGitHubRateLimit(env).catch(() => undefined); + // and being dead-lettered (the surviving event-loss path). (#audit-rate-headroom) Scoped to THIS job's own + // bucket (#audit-rate-scoping) so an unrelated installation's exhaustion never delays this job's retry. + const resetAt = await shouldWaitForGitHubRateLimit(env, LOW_REST_RATE_LIMIT_REMAINING, githubRateLimitAdmissionKeyForJob(message.body) ?? undefined).catch(() => undefined); if (resetAt) message.retry({ delaySeconds: delayUntil(resetAt) }); else message.retry(); } diff --git a/src/queue/dlq.ts b/src/queue/dlq.ts index 390d933d9b..df049ffce8 100644 --- a/src/queue/dlq.ts +++ b/src/queue/dlq.ts @@ -1,6 +1,7 @@ import { getWebhookEvent, recordAuditEvent } from "../db/repositories"; -import { delayUntil, shouldWaitForGitHubRateLimit } from "../github/rate-limit"; +import { delayUntil, LOW_REST_RATE_LIMIT_REMAINING, shouldWaitForGitHubRateLimit } from "../github/rate-limit"; import { incr } from "../selfhost/metrics"; +import { githubRateLimitAdmissionKeyForJob } from "../selfhost/queue-common"; import type { JobMessage, JsonValue } from "../types"; const DLQ_DEAD_LETTERED_METRIC = "gittensory_dlq_dead_lettered_total"; @@ -49,8 +50,10 @@ export async function processDlqBatch(batch: MessageBatch, env: Env, const event = await getWebhookEvent(env, webhook.deliveryId).catch(() => null); if (event?.status !== "processed") { // If the webhook dead-lettered because the shared GitHub REST budget was exhausted, re-drive it AFTER the - // reset (retry-until-recovered) rather than immediately re-failing it. (#audit-rate-headroom) - const resetAt = await shouldWaitForGitHubRateLimit(env).catch(() => undefined); + // reset (retry-until-recovered) rather than immediately re-failing it. (#audit-rate-headroom) Scoped to + // THIS webhook's own installation bucket (#audit-rate-scoping) so an unrelated installation's exhaustion + // never delays this re-drive. + const resetAt = await shouldWaitForGitHubRateLimit(env, LOW_REST_RATE_LIMIT_REMAINING, githubRateLimitAdmissionKeyForJob(webhook) ?? undefined).catch(() => undefined); const options = resetAt ? { delaySeconds: delayUntil(resetAt) } : undefined; const queue = env.WEBHOOKS; if (queue) { diff --git a/src/queue/processors.ts b/src/queue/processors.ts index 1c6968eae6..10c69135a1 100644 --- a/src/queue/processors.ts +++ b/src/queue/processors.ts @@ -1685,9 +1685,12 @@ async function sweepRepoRegate( // Reserve installation rate-limit headroom for real webhook traffic (#audit-rate-headroom): with the shared REST // budget at/below the maintenance floor, defer the WHOLE sweep until the reset rather than fanning out per-PR // jobs that would each have to defer. Webhooks never pre-yield, so this hands the remaining budget to them. + // Scoped to THIS repo's own installation bucket (#audit-rate-scoping) — an unrelated installation's or the + // shared public token's budget must never defer (or wrongly clear) this repo's own sweep. const sweepRateResetAt = await shouldWaitForGitHubRateLimit( env, MAINTENANCE_RESERVED_HEADROOM, + typeof repo?.installationId === "number" ? githubRateLimitAdmissionKeyForInstallation(repo.installationId) : undefined, ); if (sweepRateResetAt) { await env.JOBS.send( @@ -2024,9 +2027,12 @@ async function regatePullRequest( // repair enqueue) is current-HEAD contributor-PR-review work and gets the SAME low floor a fresh webhook // gets — it must never be treated as background maintenance and parked behind it. Mirrors the SAME // reclassification githubRateLimitAdmissionTargetForJob applies at the queue-admission layer. + // Scoped to THIS installation's own bucket (#audit-rate-scoping) — an unrelated installation's or the shared + // public token's budget must never defer (or wrongly clear) this PR's own re-gate. const rateResetAt = await shouldWaitForGitHubRateLimit( env, isScheduledRegateSweepJob(deliveryId) ? MAINTENANCE_RESERVED_HEADROOM : LOW_REST_RATE_LIMIT_REMAINING, + githubRateLimitAdmissionKeyForInstallation(installationId), ); if (rateResetAt) { await env.JOBS.send( diff --git a/src/upstream/commit.ts b/src/upstream/commit.ts index 5d7d7d00f9..bba6bf1035 100644 --- a/src/upstream/commit.ts +++ b/src/upstream/commit.ts @@ -1,5 +1,5 @@ import { githubRateLimitAdmissionKeyForPublicToken, timeoutFetch } from "../github/client"; -import { shouldWaitForGitHubRateLimit } from "../github/rate-limit"; +import { LOW_REST_RATE_LIMIT_REMAINING, shouldWaitForGitHubRateLimit } from "../github/rate-limit"; function upstreamCommitHeaders(token: string | undefined): Record { return { @@ -37,7 +37,9 @@ export async function resolveUpstreamCommitSha( // Budget-gate the NETWORK read only: a cached resolve is still served for free even under pressure; a fresh // read is skipped (→ synthetic non-OK → null → caller falls back to the mutable ref) when the REST budget is // at/below the low-water floor. This callback runs ONLY on a cache miss, so it never suppresses a cache hit. - githubSkipNetworkWhen: () => shouldWaitForGitHubRateLimit(env).then(Boolean), + // Scoped to the public-token bucket (#audit-rate-scoping) — the same bucket the read below actually draws + // from — so an installation's unrelated budget can never mask or falsely trip this gate. + githubSkipNetworkWhen: () => shouldWaitForGitHubRateLimit(env, LOW_REST_RATE_LIMIT_REMAINING, githubRateLimitAdmissionKeyForPublicToken()).then(Boolean), // Give this bare-commit read the same stable, distinctly-scoped cache-key identity every installation-token // read gets, so the shared public token rotating doesn't reset this cacheable "commit" class either (#2538). githubRateLimitAdmissionKey: githubRateLimitAdmissionKeyForPublicToken(), diff --git a/test/unit/backfill-file-hydration-scoping.test.ts b/test/unit/backfill-file-hydration-scoping.test.ts index 36f521bde4..bd965c6860 100644 --- a/test/unit/backfill-file-hydration-scoping.test.ts +++ b/test/unit/backfill-file-hydration-scoping.test.ts @@ -326,8 +326,11 @@ describe("GitHub PR file hydration scoping (#audit-rate-headroom)", () => { // Below HISTORICAL_BACKFILL_RESERVED_HEADROOM (300) but above MAINTENANCE_RESERVED_HEADROOM (150) and // LOW_REST_RATE_LIMIT_REMAINING (75) — healthy enough for the segment's own entry check and for current-PR // convergence, but not for the least-urgent historical hydration path. + // Registry-only repo (no installation), so tokenForRepo resolves to the shared public token and + // repoAdmissionKeyForToken scopes to that bucket (#audit-rate-scoping). await recordGitHubRateLimitObservation(env, { repoFullName: "JSONbored/gittensory", + admissionKey: "public-token", resource: "rest", path: "/pulls", statusCode: 200, diff --git a/test/unit/backfill.test.ts b/test/unit/backfill.test.ts index 5cf3824405..726006dbea 100644 --- a/test/unit/backfill.test.ts +++ b/test/unit/backfill.test.ts @@ -3039,6 +3039,9 @@ describe("GitHub backfill", () => { }); await recordGitHubRateLimitObservation(env, { repoFullName: "JSONbored/gittensory", + // Registry-only repo (no installation), so tokenForRepo resolves to the shared public token and + // repoAdmissionKeyForToken scopes to that bucket (#audit-rate-scoping). + admissionKey: "public-token", resource: "rest", path: "/issues", statusCode: 200, @@ -3700,6 +3703,9 @@ describe("GitHub backfill", () => { }); await recordGitHubRateLimitObservation(env, { repoFullName: "JSONbored/gittensory", + // Registry-only repo (no installation), so tokenForRepo resolves to the shared public token and + // repoAdmissionKeyForToken scopes to that bucket (#audit-rate-scoping). + admissionKey: "public-token", resource: "rest", path: "/labels", statusCode: 200, @@ -3733,6 +3739,9 @@ describe("GitHub backfill", () => { await seedRegisteredRepo(freshWaitEnv); await recordGitHubRateLimitObservation(freshWaitEnv, { repoFullName: "JSONbored/gittensory", + // Registry-only repo (no installation), so tokenForRepo resolves to the shared public token and + // repoAdmissionKeyForToken scopes to that bucket (#audit-rate-scoping). + admissionKey: "public-token", resource: "rest", path: "/labels", statusCode: 200, @@ -3802,6 +3811,9 @@ describe("GitHub backfill", () => { await seedRegisteredRepo(env); await recordGitHubRateLimitObservation(env, { repoFullName: "JSONbored/gittensory", + // Registry-only repo (no installation), so tokenForRepo resolves to the shared public token and + // repoAdmissionKeyForToken scopes to that bucket (#audit-rate-scoping). + admissionKey: "public-token", resource: "rest", path: "/pulls/1/files", statusCode: 200, @@ -3831,6 +3843,9 @@ describe("GitHub backfill", () => { await seedRegisteredRepo(env); await recordGitHubRateLimitObservation(env, { repoFullName: "JSONbored/gittensory", + // Registry-only repo (no installation), so tokenForRepo resolves to the shared public token and + // repoAdmissionKeyForToken scopes to that bucket (#audit-rate-scoping). + admissionKey: "public-token", resource: "rest", path: "/pulls", statusCode: 200, diff --git a/test/unit/index.test.ts b/test/unit/index.test.ts index 74d670c777..4e2cf3bf23 100644 --- a/test/unit/index.test.ts +++ b/test/unit/index.test.ts @@ -167,7 +167,8 @@ describe("worker entrypoint", () => { vi.useFakeTimers({ toFake: ["Date"] }); vi.setSystemTime(new Date("2026-06-24T12:00:00.000Z")); const env = createTestEnv(); - await recordGitHubRateLimitObservation(env, { repoFullName: "owner/repo", resource: "rest", path: "/x", statusCode: 200, limitValue: 5000, remaining: 120, resetAt: "2026-06-24T12:10:00.000Z", observedAt: "2026-06-24T12:00:00.000Z" }); + // Scoped to this job's own installation bucket (#audit-rate-scoping) — installationId 123 below. + await recordGitHubRateLimitObservation(env, { repoFullName: "owner/repo", admissionKey: "installation:123", resource: "rest", path: "/x", statusCode: 200, limitValue: 5000, remaining: 120, resetAt: "2026-06-24T12:10:00.000Z", observedAt: "2026-06-24T12:00:00.000Z" }); const acked: string[] = []; const retries: Array<{ delaySeconds?: number } | undefined> = []; const requeued: Array<{ message: import("../../src/types").JobMessage; delaySeconds?: number }> = []; diff --git a/test/unit/queue.test.ts b/test/unit/queue.test.ts index 702d222aa4..9f6c6d5478 100644 --- a/test/unit/queue.test.ts +++ b/test/unit/queue.test.ts @@ -7502,8 +7502,9 @@ describe("queue processors", () => { await upsertRepositorySettings(env, { repoFullName: "owner/agent-repo", autonomy: { merge: "auto" } }); await upsertPullRequestFromGitHub(env, "owner/agent-repo", { number: 7, title: "PR7", state: "open", user: { login: "c" }, head: { sha: "a7" }, labels: [], body: "" }); vi.setSystemTime(new Date("2026-05-28T02:00:00.000Z")); - // Low REST budget (10 ≤ 150 maintenance floor) with a future reset → maintenance must yield. - await repositoriesModule.recordGitHubRateLimitObservation(env, { repoFullName: "owner/agent-repo", resource: "rest", path: "/x", statusCode: 200, limitValue: 5000, remaining: 10, resetAt: "2026-05-28T02:30:00.000Z", observedAt: "2026-05-28T02:00:00.000Z" }); + // Low REST budget (10 ≤ 150 maintenance floor) with a future reset → maintenance must yield. Scoped to this + // repo's own installation bucket (#audit-rate-scoping) — the sweep now checks that bucket specifically. + await repositoriesModule.recordGitHubRateLimitObservation(env, { repoFullName: "owner/agent-repo", admissionKey: "installation:9200", resource: "rest", path: "/x", statusCode: 200, limitValue: 5000, remaining: 10, resetAt: "2026-05-28T02:30:00.000Z", observedAt: "2026-05-28T02:00:00.000Z" }); await processJob(env, { type: "agent-regate-sweep", requestedBy: "test", repoFullName: "owner/agent-repo" }); @@ -7629,7 +7630,8 @@ describe("queue processors", () => { const sent: import("../../src/types").JobMessage[] = []; const env = createTestEnv({ JOBS: { async send(m: import("../../src/types").JobMessage) { sent.push(m); } } as unknown as Queue }); vi.setSystemTime(new Date("2026-05-28T02:00:00.000Z")); - await repositoriesModule.recordGitHubRateLimitObservation(env, { repoFullName: "owner/agent-repo", resource: "rest", path: "/x", statusCode: 200, limitValue: 5000, remaining: 10, resetAt: "2026-05-28T02:30:00.000Z", observedAt: "2026-05-28T02:00:00.000Z" }); + // Scoped to this job's own installation bucket (#audit-rate-scoping) — installationId 9200 below. + await repositoriesModule.recordGitHubRateLimitObservation(env, { repoFullName: "owner/agent-repo", admissionKey: "installation:9200", resource: "rest", path: "/x", statusCode: 200, limitValue: 5000, remaining: 10, resetAt: "2026-05-28T02:30:00.000Z", observedAt: "2026-05-28T02:00:00.000Z" }); const stamp = vi.spyOn(repositoriesModule, "markPullRequestsRegated"); await processJob(env, { type: "agent-regate-pr", deliveryId: "regate-sweep:owner/agent-repo#7", repoFullName: "owner/agent-repo", prNumber: 7, installationId: 9200 }); @@ -7644,8 +7646,9 @@ describe("queue processors", () => { const env = createTestEnv({ JOBS: { async send(m: import("../../src/types").JobMessage) { sent.push(m); } } as unknown as Queue }); vi.setSystemTime(new Date("2026-05-28T02:00:00.000Z")); // 100 remaining sits BELOW the 150 maintenance floor but ABOVE the 75 live floor -- isScheduledRegateSweepJob - // must route this "regate-sweep:"-prefixed job to the higher (150) floor, so it still defers here. - await repositoriesModule.recordGitHubRateLimitObservation(env, { repoFullName: "owner/agent-repo", resource: "rest", path: "/x", statusCode: 200, limitValue: 5000, remaining: 100, resetAt: "2026-05-28T02:30:00.000Z", observedAt: "2026-05-28T02:00:00.000Z" }); + // must route this "regate-sweep:"-prefixed job to the higher (150) floor, so it still defers here. Scoped to + // this job's own installation bucket (#audit-rate-scoping) — installationId 9200 below. + await repositoriesModule.recordGitHubRateLimitObservation(env, { repoFullName: "owner/agent-repo", admissionKey: "installation:9200", resource: "rest", path: "/x", statusCode: 200, limitValue: 5000, remaining: 100, resetAt: "2026-05-28T02:30:00.000Z", observedAt: "2026-05-28T02:00:00.000Z" }); const stamp = vi.spyOn(repositoriesModule, "markPullRequestsRegated"); await processJob(env, { type: "agent-regate-pr", deliveryId: "regate-sweep:owner/agent-repo#7", repoFullName: "owner/agent-repo", prNumber: 7, installationId: 9200 }); @@ -7659,11 +7662,12 @@ describe("queue processors", () => { const sent: import("../../src/types").JobMessage[] = []; const env = createTestEnv({ JOBS: { async send(m: import("../../src/types").JobMessage) { sent.push(m); } } as unknown as Queue }); vi.setSystemTime(new Date("2026-05-28T02:00:00.000Z")); - // Same 100-remaining observation as the sibling "regate-sweep:" test above, but this deliveryId does NOT carry - // the "regate-sweep:" prefix (e.g. a repair-priority fan-out, or a real webhook-triggered re-review), so - // isScheduledRegateSweepJob is false and shouldWaitForGitHubRateLimit is called with the lower 75 floor: - // 100 > 75, so this job proceeds instead of deferring. - await repositoriesModule.recordGitHubRateLimitObservation(env, { repoFullName: "owner/agent-repo", resource: "rest", path: "/x", statusCode: 200, limitValue: 5000, remaining: 100, resetAt: "2026-05-28T02:30:00.000Z", observedAt: "2026-05-28T02:00:00.000Z" }); + // Same 100-remaining observation as the sibling "regate-sweep:" test above (scoped to this job's own + // installation:9200 bucket, #audit-rate-scoping), but this deliveryId does NOT carry the "regate-sweep:" + // prefix (e.g. a repair-priority fan-out, or a real webhook-triggered re-review), so isScheduledRegateSweepJob + // is false and shouldWaitForGitHubRateLimit is called with the lower 75 floor: 100 > 75, so this job proceeds + // instead of deferring. + await repositoriesModule.recordGitHubRateLimitObservation(env, { repoFullName: "owner/agent-repo", admissionKey: "installation:9200", resource: "rest", path: "/x", statusCode: 200, limitValue: 5000, remaining: 100, resetAt: "2026-05-28T02:30:00.000Z", observedAt: "2026-05-28T02:00:00.000Z" }); // No stored PR row for prNumber 7 -- reReviewStoredPullRequest reaches its `getPullRequest` read (proving the // rate-limit gate did not short-circuit it) and then returns immediately with no re-enqueue, since there is diff --git a/test/unit/rate-limit.test.ts b/test/unit/rate-limit.test.ts index 68fbe6b171..5c84ca70dc 100644 --- a/test/unit/rate-limit.test.ts +++ b/test/unit/rate-limit.test.ts @@ -7,8 +7,18 @@ const NOW = "2026-06-24T12:00:00.000Z"; const nowMs = Date.parse(NOW); const inIso = (ms: number): string => new Date(nowMs + ms).toISOString(); -async function seedRest(env: ReturnType, remaining: number | null, resetAt: string | null): Promise { - await recordGitHubRateLimitObservation(env, { repoFullName: "owner/repo", resource: "rest", path: "/x", statusCode: 200, limitValue: 5000, remaining, resetAt, observedAt: NOW }); +async function seedRest(env: ReturnType, remaining: number | null, resetAt: string | null, options: { admissionKey?: string; observedAt?: string } = {}): Promise { + await recordGitHubRateLimitObservation(env, { + repoFullName: "owner/repo", + admissionKey: options.admissionKey, + resource: "rest", + path: "/x", + statusCode: 200, + limitValue: 5000, + remaining, + resetAt, + observedAt: options.observedAt ?? NOW, + }); } describe("rate-limit headroom (#audit-rate-headroom)", () => { @@ -64,6 +74,49 @@ describe("rate-limit headroom (#audit-rate-headroom)", () => { expect(await shouldWaitForGitHubRateLimit(env, MAINTENANCE_RESERVED_HEADROOM)).toBe(resetAt); // 120 <= 150 → wait expect(await shouldWaitForGitHubRateLimit(env)).toBeUndefined(); // 120 > 75 default → headroom }); + + describe("admissionKey scoping (#audit-rate-scoping)", () => { + it("REGRESSION: an exhausted OTHER bucket's more-recent observation no longer throttles a caller scoped to a healthy bucket", async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date(NOW)); + const env = createTestEnv(); + // installation:2 exhausted, observed AFTER (so it would win an unscoped "globally newest" read). + await seedRest(env, 10, inIso(3_600_000), { admissionKey: "installation:2", observedAt: inIso(1_000) }); + // installation:1 healthy, observed earlier. + await seedRest(env, 500, inIso(3_600_000), { admissionKey: "installation:1", observedAt: NOW }); + expect(await shouldWaitForGitHubRateLimit(env, LOW_REST_RATE_LIMIT_REMAINING, "installation:1")).toBeUndefined(); + }); + + it("REGRESSION: a healthy OTHER bucket's more-recent observation no longer masks this caller's own exhausted bucket", async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date(NOW)); + const env = createTestEnv(); + const resetAt = inIso(3_600_000); + // installation:1 exhausted, observed earlier. + await seedRest(env, 10, resetAt, { admissionKey: "installation:1", observedAt: NOW }); + // installation:2 healthy, observed AFTER (would win an unscoped "globally newest" read and mask installation:1's exhaustion). + await seedRest(env, 500, resetAt, { admissionKey: "installation:2", observedAt: inIso(1_000) }); + expect(await shouldWaitForGitHubRateLimit(env, LOW_REST_RATE_LIMIT_REMAINING, "installation:1")).toBe(resetAt); + }); + + it("a bucket with no observations of its own returns undefined (headroom) even when another bucket is exhausted", async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date(NOW)); + const env = createTestEnv(); + await seedRest(env, 10, inIso(3_600_000), { admissionKey: "installation:1" }); + expect(await shouldWaitForGitHubRateLimit(env, LOW_REST_RATE_LIMIT_REMAINING, "public-token")).toBeUndefined(); + }); + + it("omitting admissionKey preserves the prior globally-newest-observation behavior unchanged (byte-identical fallback)", async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date(NOW)); + const env = createTestEnv(); + await seedRest(env, 500, inIso(3_600_000), { admissionKey: "installation:1", observedAt: NOW }); + const resetAt = inIso(3_600_000); + await seedRest(env, 10, resetAt, { admissionKey: "installation:2", observedAt: inIso(1_000) }); + expect(await shouldWaitForGitHubRateLimit(env)).toBe(resetAt); // picks the globally newest row (installation:2), same as before scoping existed + }); + }); }); describe("delayUntil", () => {