fix: prevent prepublication retry starvation (#3280)

convex/publishAttempts.runtime.test.ts

@@ -7,7 +7,55 @@ import schema from "./schema";
 
 const modules = import.meta.glob("./**/*.ts");
 
-describe("publish attempt orphan recovery", () => {
+describe("publish attempt runtime recovery", () => {
+  it("reserves retry-first claims for expired attempts under sustained fresh traffic", async () => {
+    const t = convexTest(schema, modules);
+    const now = Date.now();
+    const retryAttemptId = await t.run(async (ctx) => {
+      const userId = await ctx.db.insert("users", {});
+      const insertAttempt = async (
+        idempotencyKey: string,
+        createdAt: number,
+        checkClaimExpiresAt?: number,
+      ) =>
+        await ctx.db.insert("publishAttempts", {
+          kind: "skill",
+          status: "pending_checks",
+          userId,
+          slug: idempotencyKey,
+          displayName: idempotencyKey,
+          version: "1.0.0",
+          idempotencyKey,
+          artifactFingerprint: idempotencyKey,
+          files: [],
+          checks: {
+            trufflehog: { status: "pending" },
+            clawscan: { status: "pending" },
+          },
+          checkClaimExpiresAt,
+          createdAt,
+          updatedAt: createdAt,
+          expiresAt: now + 60_000,
+        });
+
+      const retryId = await insertAttempt("expired-retry", now - 60_000, now - 1);
+      for (let index = 0; index < 30; index += 1) {
+        await insertAttempt(`fresh-${index}`, now + index);
+      }
+      return retryId;
+    });
+
+    const claimed = await t.mutation(
+      internal.publishAttempts.claimPendingPublishAttemptChecksInternal,
+      {
+        claimId: "retry-reserved-claim",
+        retryOnly: true,
+      } as never,
+    );
+
+    expect(claimed).toMatchObject({ attemptId: retryAttemptId });
+  });
+
   it("terminalizes a pending attempt after its staged version is deleted", async () => {
     const t = convexTest(schema, modules);
     const ids = await t.run(async (ctx) => {

convex/publishAttempts.test.ts

@@ -396,6 +396,57 @@ describe("publishAttempts", () => {
     expect(ctx.storage.getUrl).not.toHaveBeenCalled();
   });
 
+  it("uses reserved claims for expired pending-check retries before finalization work", async () => {
+    const previousToken = process.env.SECURITY_SCAN_WORKER_TOKEN;
+    process.env.SECURITY_SCAN_WORKER_TOKEN = "worker-token";
+    const retry = {
+      attemptId: "publishAttempts:retry",
+      status: "pending_checks",
+      claimId: "claim-1",
+      kind: "skill",
+      userId: "users:publisher",
+      slug: "retry-skill",
+      displayName: "Retry Skill",
+      version: "1.0.0",
+      artifactFingerprint: "fingerprint",
+      files: [],
+      checkClaimExpiresAt: Date.now() + 60_000,
+      createdAt: Date.now(),
+    };
+    const ctx = {
+      runMutation: vi.fn(async (ref: Parameters<typeof getFunctionName>[0], _args?: unknown) => {
+        const name = getFunctionName(ref);
+        return name === "publishAttempts:claimPendingPublishAttemptChecksInternal"
+          ? retry
+          : { ...retry, attemptId: "publishAttempts:ready", status: "ready_to_finalize" };
+      }),
+      storage: {
+        getUrl: vi.fn(),
+      },
+    };
+
+    try {
+      await expect(
+        claimPrePublicationChecksHandler(ctx, {
+          token: "worker-token",
+          preferRetry: true,
+        }),
+      ).resolves.toMatchObject({
+        attemptId: "publishAttempts:retry",
+        status: "pending_checks",
+      });
+    } finally {
+      if (previousToken === undefined) delete process.env.SECURITY_SCAN_WORKER_TOKEN;
+      else process.env.SECURITY_SCAN_WORKER_TOKEN = previousToken;
+    }
+
+    expect(ctx.runMutation).toHaveBeenCalledTimes(1);
+    expect(
+      getFunctionName(ctx.runMutation.mock.calls[0]?.[0] as Parameters<typeof getFunctionName>[0]),
+    ).toBe("publishAttempts:claimPendingPublishAttemptChecksInternal");
+    expect(ctx.runMutation.mock.calls[0]?.[1]).toMatchObject({ retryOnly: true });
+  });
+
   it("lets targeted pending attempts fall through the ready-finalization lookup", async () => {
     const ctx = {
       db: {

convex/publishAttempts.ts

@@ -712,23 +712,36 @@ export const claimPendingPublishAttemptChecksInternal = internalMutation({
     claimId: v.string(),
     attemptId: v.optional(v.id("publishAttempts")),
     kind: v.optional(v.union(v.literal("skill"), v.literal("package"))),
+    retryOnly: v.optional(v.boolean()),
     slug: v.optional(v.string()),
     version: v.optional(v.string()),
   },
   handler: async (ctx, args) => {
     const now = Date.now();
     const targetedAttempt = args.attemptId ? await ctx.db.get(args.attemptId) : null;
-    const candidates = args.attemptId
-      ? targetedAttempt
-        ? [targetedAttempt]
-        : []
-      : await ctx.db
-          .query("publishAttempts")
-          .withIndex("by_status_check_claim_expires_at_created", (q) =>
-            q.eq("status", "pending_checks"),
-          )
-          .order("asc")
-          .take(25);
+    let candidates: Doc<"publishAttempts">[];
+    if (args.attemptId) {
+      candidates = targetedAttempt ? [targetedAttempt] : [];
+    } else if (args.retryOnly) {
+      candidates = await ctx.db
+        .query("publishAttempts")
+        .withIndex("by_status_check_claim_expires_at_created", (q) =>
+          q
+            .eq("status", "pending_checks")
+            .gte("checkClaimExpiresAt", 0)
+            .lte("checkClaimExpiresAt", now),
+        )
+        .order("asc")
+        .take(25);
+    } else {
+      candidates = await ctx.db
+        .query("publishAttempts")
+        .withIndex("by_status_check_claim_expires_at_created", (q) =>
+          q.eq("status", "pending_checks"),
+        )
+        .order("asc")
+        .take(25);
+    }
 
     for (const attempt of candidates) {
       if (attempt.status !== "pending_checks") {
@@ -1177,6 +1190,7 @@ export const claimPrePublicationChecks: ReturnType<typeof action> = action({
     token: v.string(),
     attemptId: v.optional(v.id("publishAttempts")),
     kind: v.optional(v.union(v.literal("skill"), v.literal("package"))),
+    preferRetry: v.optional(v.boolean()),
     slug: v.optional(v.string()),
     version: v.optional(v.string()),
   },
@@ -1190,10 +1204,17 @@ export const claimPrePublicationChecks: ReturnType<typeof action> = action({
       slug: args.slug,
       version: args.version,
     };
-    const claimed = ((await ctx.runMutation(
-      internal.publishAttempts.claimReadyPublishAttemptFinalizationRetryInternal,
-      claimArgs,
-    )) ??
+    const reservedRetry = args.preferRetry
+      ? await ctx.runMutation(internal.publishAttempts.claimPendingPublishAttemptChecksInternal, {
+          ...claimArgs,
+          retryOnly: true,
+        })
+      : null;
+    const claimed = (reservedRetry ??
+      (await ctx.runMutation(
+        internal.publishAttempts.claimReadyPublishAttemptFinalizationRetryInternal,
+        claimArgs,
+      )) ??
       (await ctx.runMutation(
         internal.publishAttempts.claimPendingPublishAttemptChecksInternal,
         claimArgs,

scripts/security/run-prepublication-worker.test.ts

@@ -375,6 +375,29 @@ describe("pre-publication worker", () => {
     expect(client.action).toHaveBeenCalledTimes(2);
   });
 
+  it("reserves one batch claim for expired retries, including filtered batches", async () => {
+    const client = {
+      action: vi.fn().mockResolvedValue(null),
+    };
+
+    await expect(
+      claimPrePublicationBatch(client, "worker-token", 4, { kind: "skill" }),
+    ).resolves.toEqual({
+      attempts: [],
+      claimFailures: 0,
+    });
+    expect(client.action.mock.calls.map(([, args]) => args.preferRetry)).toEqual([
+      true,
+      undefined,
+      undefined,
+      undefined,
+    ]);
+    expect(client.action).toHaveBeenCalledWith(
+      expect.anything(),
+      expect.objectContaining({ kind: "skill", preferRetry: true }),
+    );
+  });
+
   it("fails the worker batch when claims fail without claiming work", async () => {
     const client = {
       action: vi

scripts/security/run-prepublication-worker.ts

@@ -791,10 +791,12 @@ export async function claimPrePublicationAttempt(
   client: PrePublicationWorkerClient,
   token: string,
   filters: PrePublicationClaimFilters = {},
+  preferRetry = false,
 ) {
   return (await client.action(api.publishAttempts.claimPrePublicationChecks, {
     token,
     ...filters,
+    ...(preferRetry ? { preferRetry: true } : {}),
   })) as ClaimedPrePublicationAttempt | null;
 }
 
@@ -805,7 +807,9 @@ export async function claimPrePublicationBatch(
   filters: PrePublicationClaimFilters = {},
 ) {
   const claims = await Promise.allSettled(
-    Array.from({ length: limit }, () => claimPrePublicationAttempt(client, token, filters)),
+    Array.from({ length: limit }, (_, index) =>
+      claimPrePublicationAttempt(client, token, filters, index === 0),
+    ),
   );
   const attempts: ClaimedPrePublicationAttempt[] = [];
   const failures: unknown[] = [];