DX-2969: retry blob credentials and delete while a new bucket provisions

src/commands/blob/credentials.ts

@@ -3,17 +3,35 @@ import { resolveAuth } from "../../auth.js";
 import { HttpError, request } from "../../client.js";
 import { printJSON } from "../../output.js";
 import type { BlobBucket, BlobS3Credentials } from "../../types.js";
+import {
+  isFreshlyCreated,
+  PROVISIONING_MAX_RETRIES,
+  PROVISIONING_RETRY_DELAY_MS,
+  sleep,
+} from "./retry.js";
+import type { Sleep } from "./retry.js";
 
 const BLOB_CREDENTIALS_URL = "https://blob.upstash.io/v1/credentials";
 const RETRYABLE_STATUSES = new Set([429, 503]);
 const DEFAULT_RETRY_DELAY_MS = 2000;
 const MAX_RETRY_DELAY_MS = 10000;
 const MAX_RETRIES = 3;
 
-type Sleep = (ms: number) => Promise<void>;
+export interface FetchBlobCredentialsOptions {
+  /**
+   * How many times a 401 may be retried. A bucket that was created moments ago
+   * can exist in the Developer API before the Blob worker has provisioned it,
+   * and the worker answers 401 during that window.
+   */
+  unauthorizedRetries?: number;
+}
 
-function sleep(ms: number): Promise<void> {
-  return new Promise((resolve) => setTimeout(resolve, ms));
+function unauthorizedMessage(retriesUsed: number): string {
+  if (retriesUsed > 0) {
+    const waited = Math.round((retriesUsed * PROVISIONING_RETRY_DELAY_MS) / 1000);
+    return `Blob bucket token was rejected after waiting ${waited}s for the bucket to finish provisioning; retry in a minute or check the bucket in the console`;
+  }
+  return "Blob bucket token was rejected; if the bucket was created moments ago it may still be provisioning, wait a few seconds and retry";
 }
 
 function parseErrorMessage(text: string, status: number): string {
@@ -86,8 +104,13 @@ function validateCredentials(data: unknown): BlobS3Credentials {
 export async function fetchBlobCredentials(
   token: string,
   pause: Sleep = sleep,
+  options: FetchBlobCredentialsOptions = {},
 ): Promise<BlobS3Credentials> {
-  for (let attempt = 0; attempt <= MAX_RETRIES; attempt += 1) {
+  const unauthorizedRetries = options.unauthorizedRetries ?? 0;
+  let throttled = 0;
+  let unauthorized = 0;
+
+  for (;;) {
     const response = await fetch(BLOB_CREDENTIALS_URL, {
       method: "POST",
       headers: {
@@ -108,34 +131,50 @@ export async function fetchBlobCredentials(
     }
 
     if (response.status === 401) {
-      throw new HttpError("Blob bucket token was rejected", response.status);
+      if (unauthorized < unauthorizedRetries) {
+        unauthorized += 1;
+        await pause(PROVISIONING_RETRY_DELAY_MS);
+        continue;
+      }
+      throw new HttpError(unauthorizedMessage(unauthorized), response.status);
     }
 
-    if (RETRYABLE_STATUSES.has(response.status) && attempt < MAX_RETRIES) {
+    if (RETRYABLE_STATUSES.has(response.status) && throttled < MAX_RETRIES) {
+      throttled += 1;
       await pause(getRetryDelayMs(response.headers.get("Retry-After")));
       continue;
     }
 
     throw new HttpError(parseErrorMessage(text, response.status), response.status);
   }
+}
 
-  throw new Error("Blob credentials request failed after retries");
+interface BucketTokenSource {
+  token: string;
+  /** Retry 401s for buckets that may still be provisioning; 0 for ambient tokens. */
+  unauthorizedRetries: number;
 }
 
-function resolveBucketToken(flags: { bucketId?: string }, command: Command): Promise<string> {
+function resolveBucketToken(
+  flags: { bucketId?: string },
+  command: Command,
+): Promise<BucketTokenSource> {
   if (flags.bucketId) {
     const auth = resolveAuth(command);
     return request<BlobBucket>(auth, "GET", `/v2/blob/bucket/${flags.bucketId}`).then((bucket) => {
       if (typeof bucket.token === "string" && bucket.token.length > 0) {
-        return bucket.token;
+        return {
+          token: bucket.token,
+          unauthorizedRetries: isFreshlyCreated(bucket.creation_time) ? PROVISIONING_MAX_RETRIES : 0,
+        };
       }
       throw new Error(`Blob bucket ${flags.bucketId} did not return a current token`);
     });
   }
 
   const token = process.env.UPSTASH_BLOB_TOKEN;
   if (typeof token === "string" && token.length > 0) {
-    return Promise.resolve(token);
+    return Promise.resolve({ token, unauthorizedRetries: 0 });
   }
 
   return Promise.reject(
@@ -152,9 +191,18 @@ export function registerBlobCredentials(blob: Command): void {
       "Get temporary S3 credentials for a Blob bucket; expiresAt is the credential expiry",
     )
     .option("--bucket-id <id>", "Blob bucket ID")
+    .addHelpText(
+      "after",
+      `
+With --bucket-id, a bucket created in the last few minutes is polled for up
+to ~30s until provisioning finishes, so it is safe to run right after create.
+`,
+    )
     .action(async (flags: { bucketId?: string }, command: Command) => {
-      const token = await resolveBucketToken(flags, command);
-      const credentials = await fetchBlobCredentials(token);
+      const source = await resolveBucketToken(flags, command);
+      const credentials = await fetchBlobCredentials(source.token, sleep, {
+        unauthorizedRetries: source.unauthorizedRetries,
+      });
       printJSON(credentials);
     });
 }

src/commands/blob/delete.ts

@@ -1,7 +1,41 @@
 import { Command } from "commander";
 import { resolveAuth } from "../../auth.js";
-import { request } from "../../client.js";
+import { HttpError, request } from "../../client.js";
 import { printJSON } from "../../output.js";
+import { sleep } from "./retry.js";
+import type { Sleep } from "./retry.js";
+import type { Auth } from "../../auth.js";
+
+const SERVER_ERROR_RETRY_DELAY_MS = 3000;
+const SERVER_ERROR_MAX_RETRIES = 5;
+
+/**
+ * Deleting a bucket moments after creating it can fail with a 5xx while the
+ * backend is still provisioning it. Retry briefly; a 404 after an earlier
+ * attempt means that attempt actually went through.
+ */
+export async function deleteBlobBucket(
+  auth: Auth,
+  bucketId: string,
+  pause: Sleep = sleep,
+): Promise<void> {
+  let retries = 0;
+  for (;;) {
+    try {
+      await request(auth, "DELETE", `/v2/blob/bucket/${bucketId}`);
+      return;
+    } catch (error) {
+      if (!(error instanceof HttpError)) throw error;
+      if (error.status === 404 && retries > 0) return;
+      if (error.status >= 500 && retries < SERVER_ERROR_MAX_RETRIES) {
+        retries += 1;
+        await pause(SERVER_ERROR_RETRY_DELAY_MS);
+        continue;
+      }
+      throw error;
+    }
+  }
+}
 
 export function registerBlobDelete(blob: Command): void {
   blob
@@ -15,7 +49,7 @@ export function registerBlobDelete(blob: Command): void {
         return;
       }
       const auth = resolveAuth(command);
-      await request(auth, "DELETE", `/v2/blob/bucket/${flags.bucketId}`);
+      await deleteBlobBucket(auth, flags.bucketId);
       printJSON({ deleted: true, bucket_id: flags.bucketId });
     });
 }

src/commands/blob/retry.ts

@@ -0,0 +1,15 @@
+export type Sleep = (ms: number) => Promise<void>;
+
+export function sleep(ms: number): Promise<void> {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}
+
+/** Buckets younger than this may still be provisioning on the Blob worker. */
+export const PROVISIONING_WINDOW_SECONDS = 5 * 60;
+export const PROVISIONING_RETRY_DELAY_MS = 3000;
+export const PROVISIONING_MAX_RETRIES = 10;
+
+export function isFreshlyCreated(creationTime: number | undefined, nowSeconds = Date.now() / 1000): boolean {
+  if (typeof creationTime !== "number" || !Number.isFinite(creationTime)) return true;
+  return nowSeconds - creationTime < PROVISIONING_WINDOW_SECONDS;
+}

tests/unit/blob.test.ts

@@ -1,6 +1,8 @@
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import { createBlobProgram, runCommand } from "../helpers/program.js";
 import { fetchBlobCredentials } from "../../src/commands/blob/credentials.js";
+import { deleteBlobBucket } from "../../src/commands/blob/delete.js";
+import { isFreshlyCreated } from "../../src/commands/blob/retry.js";
 import type { BlobBucket, BlobS3Credentials } from "../../src/types.js";
 
 const originalEnv = { ...process.env };
@@ -317,6 +319,93 @@ describe("blob credentials command", () => {
     ).rejects.toThrow(/rejected/);
   });
 
+  it("401 hint mentions provisioning when no retries were allowed", async () => {
+    vi.spyOn(globalThis, "fetch").mockResolvedValue(
+      new Response('{"error":"unauthorized"}', { status: 401 }),
+    );
+
+    await expect(fetchBlobCredentials("token", async () => {})).rejects.toThrow(
+      /still be provisioning/,
+    );
+  });
+
+  it("retries 401 while a fresh bucket is provisioning, then succeeds", async () => {
+    const credentials = makeCredentials();
+    const delays: number[] = [];
+    vi.spyOn(globalThis, "fetch")
+      .mockResolvedValueOnce(new Response('{"error":"unauthorized"}', { status: 401 }))
+      .mockResolvedValueOnce(new Response('{"error":"unauthorized"}', { status: 401 }))
+      .mockResolvedValueOnce(new Response(JSON.stringify(credentials), { status: 200 }));
+
+    const result = await fetchBlobCredentials(
+      "bucket-token",
+      async (ms) => {
+        delays.push(ms);
+      },
+      { unauthorizedRetries: 10 },
+    );
+
+    expect(result).toEqual(credentials);
+    expect(delays).toEqual([3000, 3000]);
+  });
+
+  it("gives up on 401 after the provisioning retry budget and says how long it waited", async () => {
+    const fetchSpy = vi.spyOn(globalThis, "fetch").mockImplementation(async () =>
+      new Response('{"error":"unauthorized"}', { status: 401 }),
+    );
+
+    await expect(
+      fetchBlobCredentials("bucket-token", async () => {}, { unauthorizedRetries: 2 }),
+    ).rejects.toThrow(/rejected after waiting 6s/);
+    expect(fetchSpy).toHaveBeenCalledTimes(3);
+  });
+
+  it("401 retries do not consume the throttle retry budget", async () => {
+    const credentials = makeCredentials();
+    vi.spyOn(globalThis, "fetch")
+      .mockResolvedValueOnce(new Response("", { status: 401 }))
+      .mockResolvedValueOnce(new Response("", { status: 503 }))
+      .mockResolvedValueOnce(new Response("", { status: 401 }))
+      .mockResolvedValueOnce(new Response("", { status: 503 }))
+      .mockResolvedValueOnce(new Response("", { status: 503 }))
+      .mockResolvedValueOnce(new Response(JSON.stringify(credentials), { status: 200 }));
+
+    const result = await fetchBlobCredentials("t", async () => {}, { unauthorizedRetries: 2 });
+    expect(result).toEqual(credentials);
+  });
+
+  it("by bucket id does not retry 401 for a bucket that is not freshly created", async () => {
+    const bucket = makeBucket({ creation_time: 1 });
+    const fetchSpy = vi
+      .spyOn(globalThis, "fetch")
+      .mockResolvedValueOnce(new Response(JSON.stringify(bucket), { status: 200 }))
+      .mockResolvedValue(new Response('{"error":"unauthorized"}', { status: 401 }));
+
+    const program = await createBlobProgram();
+    await expect(
+      runCommand(program, ["blob", "credentials", "--bucket-id", "bucket_123"]),
+    ).rejects.toThrow(/rejected/);
+    expect(fetchSpy).toHaveBeenCalledTimes(2);
+  });
+
+  it("by bucket id polls 401 for a freshly created bucket", async () => {
+    vi.useFakeTimers();
+    const bucket = makeBucket({ creation_time: Math.floor(Date.now() / 1000) });
+    const credentials = makeCredentials();
+    const fetchSpy = vi
+      .spyOn(globalThis, "fetch")
+      .mockResolvedValueOnce(new Response(JSON.stringify(bucket), { status: 200 }))
+      .mockResolvedValueOnce(new Response('{"error":"unauthorized"}', { status: 401 }))
+      .mockResolvedValueOnce(new Response(JSON.stringify(credentials), { status: 200 }));
+
+    const program = await createBlobProgram();
+    const pending = runCommand(program, ["blob", "credentials", "--bucket-id", "bucket_123"]);
+    await vi.advanceTimersByTimeAsync(3000);
+
+    expect(await pending).toEqual(credentials);
+    expect(fetchSpy).toHaveBeenCalledTimes(3);
+  });
+
   it("retries 429 and 503 with retry-after or fallback delays, then succeeds", async () => {
     const credentials = makeCredentials();
     const delays: number[] = [];
@@ -349,3 +438,70 @@ describe("blob credentials command", () => {
     expect(delays).toEqual([10000, 2000, 3000]);
   });
 });
+
+describe("blob provisioning helpers", () => {
+  it("treats missing or recent creation times as freshly created", () => {
+    expect(isFreshlyCreated(undefined)).toBe(true);
+    expect(isFreshlyCreated(1000, 1000 + 60)).toBe(true);
+    expect(isFreshlyCreated(1000, 1000 + 5 * 60)).toBe(false);
+  });
+});
+
+describe("blob delete retries", () => {
+  const auth = { email: "user@example.com", apiKey: "api-key" };
+
+  it("retries 5xx while the bucket is provisioning, then succeeds", async () => {
+    const delays: number[] = [];
+    const fetchSpy = vi
+      .spyOn(globalThis, "fetch")
+      .mockResolvedValueOnce(new Response('{"error":"internal"}', { status: 500 }))
+      .mockResolvedValueOnce(new Response('"OK"', { status: 200 }));
+
+    await deleteBlobBucket(auth, "bucket_123", async (ms) => {
+      delays.push(ms);
+    });
+
+    expect(fetchSpy).toHaveBeenCalledTimes(2);
+    expect(delays).toEqual([3000]);
+  });
+
+  it("treats a 404 after a failed attempt as already deleted", async () => {
+    vi.spyOn(globalThis, "fetch")
+      .mockResolvedValueOnce(new Response('{"error":"internal"}', { status: 500 }))
+      .mockResolvedValueOnce(new Response('{"error":"resource not found"}', { status: 404 }));
+
+    await expect(deleteBlobBucket(auth, "bucket_123", async () => {})).resolves.toBeUndefined();
+  });
+
+  it("does not retry a first-attempt 404 or any 4xx", async () => {
+    const fetchSpy = vi.spyOn(globalThis, "fetch")
+      .mockResolvedValueOnce(new Response('{"error":"resource not found"}', { status: 404 }));
+    await expect(deleteBlobBucket(auth, "bucket_123", async () => {})).rejects.toThrow(/not found/);
+    expect(fetchSpy).toHaveBeenCalledTimes(1);
+
+    fetchSpy.mockResolvedValueOnce(new Response('{"error":"bad"}', { status: 400 }));
+    await expect(deleteBlobBucket(auth, "bucket_123", async () => {})).rejects.toThrow(/bad/);
+    expect(fetchSpy).toHaveBeenCalledTimes(2);
+  });
+
+  it("gives up after the 5xx retry budget", async () => {
+    const fetchSpy = vi.spyOn(globalThis, "fetch").mockImplementation(async () =>
+      new Response('{"error":"internal"}', { status: 500 }),
+    );
+    await expect(deleteBlobBucket(auth, "bucket_123", async () => {})).rejects.toThrow(/internal/);
+    expect(fetchSpy).toHaveBeenCalledTimes(6);
+  });
+
+  it("delete command surfaces the retried result", async () => {
+    vi.useFakeTimers();
+    vi.spyOn(globalThis, "fetch")
+      .mockResolvedValueOnce(new Response('{"error":"internal"}', { status: 500 }))
+      .mockResolvedValueOnce(new Response('"OK"', { status: 200 }));
+
+    const program = await createBlobProgram();
+    const pending = runCommand(program, ["blob", "delete", "--bucket-id", "bucket_123"]);
+    await vi.advanceTimersByTimeAsync(3000);
+
+    expect(await pending).toEqual({ deleted: true, bucket_id: "bucket_123" });
+  });
+});