From c843a4230fe0fbcc8f92aa5c388a8d3fa42e64e3 Mon Sep 17 00:00:00 2001
From: lukebjerring <lukebjerring@gmail.com>
Date: Fri, 11 Sep 2026 13:34:13 +0000
Subject: [PATCH] setup_builder: retry transient sticky disk requests

Co-authored-by: Codesmith Staging <codesmith-bot@users.noreply.github.com>
---
 src/setup-builder.test.ts | 57 +++++++++++++++++++++++++++++++++++++--
 src/setup_builder.ts      | 51 +++++++++++++++++++++++++----------
 2 files changed, 92 insertions(+), 16 deletions(-)

diff --git a/src/setup-builder.test.ts b/src/setup-builder.test.ts
index 197f9cc..b97ac1c 100644
--- a/src/setup-builder.test.ts
+++ b/src/setup-builder.test.ts
@@ -1,8 +1,9 @@
-import { describe, it, expect, vi, beforeEach } from "vitest";
+import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
 import * as core from "@actions/core";
 import * as fs from "fs";
 import * as setupBuilder from "./setup_builder";
-// import * as reporter from "./reporter";
+import * as reporter from "./reporter";
+import { Code, ConnectError } from "@connectrpc/connect";
 
 // Mock the modules
 vi.mock("@actions/core", () => ({
@@ -48,6 +49,58 @@ describe("setup_builder", () => {
     process.env.BLACKSMITH_VM_ID = "test-vm-id";
   });
 
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
+  describe("getStickyDisk", () => {
+    it("retries transient 503 responses", async () => {
+      vi.useFakeTimers();
+      const getStickyDisk = vi
+        .fn()
+        .mockRejectedValueOnce(
+          new ConnectError(
+            "failed to expose and mount sticky disk: unexpected status code: 503",
+            Code.Unknown,
+          ),
+        )
+        .mockResolvedValue({
+          exposeId: "expose-1",
+          diskIdentifier: "/dev/vdb",
+        });
+      vi.mocked(reporter.createBlacksmithAgentClient).mockResolvedValue({
+        up: vi.fn(),
+        getStickyDisk,
+      } as never);
+
+      const resultPromise = setupBuilder.getStickyDisk();
+      await vi.runAllTimersAsync();
+
+      await expect(resultPromise).resolves.toMatchObject({
+        expose_id: "expose-1",
+        device: "/dev/vdb",
+      });
+      expect(getStickyDisk).toHaveBeenCalledTimes(2);
+    });
+
+    it("does not retry 500 responses", async () => {
+      const getStickyDisk = vi
+        .fn()
+        .mockRejectedValue(
+          new ConnectError("unexpected status code: 500", Code.Unknown),
+        );
+      vi.mocked(reporter.createBlacksmithAgentClient).mockResolvedValue({
+        up: vi.fn(),
+        getStickyDisk,
+      } as never);
+
+      await expect(setupBuilder.getStickyDisk()).rejects.toThrow(
+        "unexpected status code: 500",
+      );
+      expect(getStickyDisk).toHaveBeenCalledOnce();
+    });
+  });
+
   describe("getNumCPUs", () => {
     it("should return the number of CPUs", async () => {
       const exec = (await import("child_process")).exec as unknown as {
diff --git a/src/setup_builder.ts b/src/setup_builder.ts
index 5c95db2..abb9d29 100644
--- a/src/setup_builder.ts
+++ b/src/setup_builder.ts
@@ -7,6 +7,7 @@ import * as reporter from "./reporter";
 import { execa } from "execa";
 import * as stateHelper from "./state-helper";
 import { BOLT_CHECK_MAX_FILE_BYTES } from "./exec-utils";
+import { Code, ConnectError } from "@connectrpc/connect";
 
 // Constants for configuration.
 const BUILDKIT_DAEMON_ADDR = "tcp://127.0.0.1:1234";
@@ -415,20 +416,42 @@ export async function getStickyDisk(options?: {
   }
   core.info(`Getting sticky disk for ${stickyDiskKey}`);
 
-  const response = await client.getStickyDisk(
-    {
-      stickyDiskKey: stickyDiskKey,
-      region: process.env.BLACKSMITH_REGION || "eu-central",
-      installationModelId: process.env.BLACKSMITH_INSTALLATION_MODEL_ID || "",
-      vmId: process.env.BLACKSMITH_VM_ID || "",
-      stickyDiskType: "dockerfile",
-      repoName: process.env.GITHUB_REPO_NAME || "",
-      stickyDiskToken: process.env.BLACKSMITH_STICKYDISK_TOKEN || "",
-    },
-    {
-      signal: options?.signal,
-    },
-  );
+  const request = {
+    stickyDiskKey,
+    region: process.env.BLACKSMITH_REGION || "eu-central",
+    installationModelId: process.env.BLACKSMITH_INSTALLATION_MODEL_ID || "",
+    vmId: process.env.BLACKSMITH_VM_ID || "",
+    stickyDiskType: "dockerfile",
+    repoName: process.env.GITHUB_REPO_NAME || "",
+    stickyDiskToken: process.env.BLACKSMITH_STICKYDISK_TOKEN || "",
+  };
+
+  let response;
+  for (let attempt = 1; ; attempt++) {
+    try {
+      response = await client.getStickyDisk(request, {
+        signal: options?.signal,
+      });
+      break;
+    } catch (error) {
+      const connectError = ConnectError.from(error);
+      const retryable =
+        connectError.code === Code.Unavailable ||
+        /unexpected status code: 50[234]\b/.test(connectError.rawMessage);
+      if (!retryable || attempt === 3 || options?.signal?.aborted) {
+        throw error;
+      }
+
+      const delayMs = attempt * 500;
+      core.warning(
+        `Sticky disk request failed (attempt ${attempt}/3), retrying in ${delayMs}ms: ${connectError.rawMessage}`,
+      );
+      await new Promise((resolve) => setTimeout(resolve, delayMs));
+      if (options?.signal?.aborted) {
+        throw error;
+      }
+    }
+  }
   return {
     expose_id: (response as { exposeId?: string }).exposeId || "",
     device: (response as { diskIdentifier?: string }).diskIdentifier || "",