feat(supervisor): optional image registry rewrite for run pods

Author: nicktrn

apps/supervisor/src/env.ts

@@ -185,6 +185,8 @@ export const Env = z
     KUBERNETES_EPHEMERAL_STORAGE_SIZE_LIMIT: z.string().default("10Gi"),
     KUBERNETES_EPHEMERAL_STORAGE_SIZE_REQUEST: z.string().default("2Gi"),
     KUBERNETES_STRIP_IMAGE_DIGEST: BoolEnv.default(false),
+    KUBERNETES_IMAGE_REGISTRY_REWRITE_FROM: z.string().optional(),
+    KUBERNETES_IMAGE_REGISTRY_REWRITE_TO: z.string().optional(),
     KUBERNETES_CPU_REQUEST_MIN_CORES: z.coerce.number().min(0).default(0),
     KUBERNETES_CPU_REQUEST_RATIO: z.coerce.number().min(0).max(1).default(0.75), // Ratio of CPU limit, so 0.75 = 75% of CPU limit
     KUBERNETES_MEMORY_REQUEST_MIN_GB: z.coerce.number().min(0).default(0),

apps/supervisor/src/workloadManager/imageRegistry.test.ts

@@ -0,0 +1,42 @@
+import { describe, expect, it } from "vitest";
+import { rewriteImageRegistry } from "./imageRegistry.js";
+
+const FROM = "123456789012.dkr.ecr.us-east-1.amazonaws.com";
+const TO = "123456789012.dkr.ecr.eu-central-1.amazonaws.com";
+
+describe("rewriteImageRegistry", () => {
+  it("rewrites the registry host and keeps the rest of the reference", () => {
+    expect(rewriteImageRegistry(`${FROM}/deployments/proj_abc:20260818.1`, FROM, TO)).toBe(
+      `${TO}/deployments/proj_abc:20260818.1`
+    );
+  });
+
+  it("preserves a digest", () => {
+    expect(rewriteImageRegistry(`${FROM}/deployments/proj_abc@sha256:abc123`, FROM, TO)).toBe(
+      `${TO}/deployments/proj_abc@sha256:abc123`
+    );
+  });
+
+  it("is a no-op unless both ends are configured", () => {
+    const ref = `${FROM}/deployments/proj_abc:tag`;
+
+    expect(rewriteImageRegistry(ref, undefined, TO)).toBe(ref);
+    expect(rewriteImageRegistry(ref, FROM, undefined)).toBe(ref);
+    expect(rewriteImageRegistry(ref, undefined, undefined)).toBe(ref);
+  });
+
+  it("leaves other registries alone", () => {
+    const ref = "ghcr.io/triggerdotdev/something:tag";
+    expect(rewriteImageRegistry(ref, FROM, TO)).toBe(ref);
+  });
+
+  it("only matches on a host boundary", () => {
+    const lookalike = `${FROM}.evil.example.com/deployments/proj_abc:tag`;
+    expect(rewriteImageRegistry(lookalike, FROM, TO)).toBe(lookalike);
+  });
+
+  it("does not rewrite a host that merely contains the source", () => {
+    const ref = `registry.example.com/${FROM}/proj_abc:tag`;
+    expect(rewriteImageRegistry(ref, FROM, TO)).toBe(ref);
+  });
+});

apps/supervisor/src/workloadManager/imageRegistry.ts

@@ -0,0 +1,15 @@
+export function rewriteImageRegistry(
+  imageRef: string,
+  from: string | undefined,
+  to: string | undefined
+): string {
+  if (!from || !to) {
+    return imageRef;
+  }
+
+  if (!imageRef.startsWith(`${from}/`)) {
+    return imageRef;
+  }
+
+  return `${to}${imageRef.slice(from.length)}`;
+}

apps/supervisor/src/workloadManager/kubernetes.ts

@@ -20,6 +20,7 @@ import {
   withBlockIoUringSeccompProfile,
   withNodeSelector,
 } from "./kubernetesPodSpec.js";
+import { rewriteImageRegistry } from "./imageRegistry.js";
 
 type ResourceQuantities = {
   [K in "cpu" | "memory" | "ephemeral-storage"]?: string;
@@ -160,7 +161,11 @@ export class KubernetesWorkloadManager implements WorkloadManager {
             containers: [
               {
                 name: "run-controller",
-                image: this.stripImageDigest(opts.image),
+                image: rewriteImageRegistry(
+                  this.stripImageDigest(opts.image),
+                  env.KUBERNETES_IMAGE_REGISTRY_REWRITE_FROM,
+                  env.KUBERNETES_IMAGE_REGISTRY_REWRITE_TO
+                ),
                 ports: [
                   {
                     containerPort: 8000,