ability to set arbitrary env vars on new runners

Author: nicktrn

apps/supervisor/src/env.ts

@@ -1,14 +1,7 @@
 import { randomUUID } from "crypto";
 import { env as stdEnv } from "std-env";
 import { z } from "zod";
-
-const BoolEnv = z.preprocess((val) => {
-  if (typeof val !== "string") {
-    return val;
-  }
-
-  return ["true", "1"].includes(val.toLowerCase().trim());
-}, z.boolean());
+import { AdditionalEnvVars, BoolEnv } from "./envUtil.js";
 
 const Env = z.object({
   // This will come from `spec.nodeName` in k8s
@@ -75,6 +68,9 @@ const Env = z.object({
 
   // Debug
   DEBUG: BoolEnv.default(false),
+
+  // Additional environment variables (CSV format)
+  RUNNER_ADDITIONAL_ENV_VARS: AdditionalEnvVars,
 });
 
 export const env = Env.parse(stdEnv);

apps/supervisor/src/envUtil.test.ts

@@ -0,0 +1,80 @@
+import { describe, it, expect } from "vitest";
+import { BoolEnv, AdditionalEnvVars } from "./envUtil.js";
+
+describe("BoolEnv", () => {
+  it("should parse string 'true' as true", () => {
+    expect(BoolEnv.parse("true")).toBe(true);
+    expect(BoolEnv.parse("TRUE")).toBe(true);
+    expect(BoolEnv.parse("True")).toBe(true);
+  });
+
+  it("should parse string '1' as true", () => {
+    expect(BoolEnv.parse("1")).toBe(true);
+  });
+
+  it("should parse string 'false' as false", () => {
+    expect(BoolEnv.parse("false")).toBe(false);
+    expect(BoolEnv.parse("FALSE")).toBe(false);
+    expect(BoolEnv.parse("False")).toBe(false);
+  });
+
+  it("should handle whitespace", () => {
+    expect(BoolEnv.parse(" true ")).toBe(true);
+    expect(BoolEnv.parse(" 1 ")).toBe(true);
+  });
+
+  it("should pass through boolean values", () => {
+    expect(BoolEnv.parse(true)).toBe(true);
+    expect(BoolEnv.parse(false)).toBe(false);
+  });
+
+  it("should return false for invalid inputs", () => {
+    expect(BoolEnv.parse("invalid")).toBe(false);
+    expect(BoolEnv.parse("")).toBe(false);
+  });
+});
+
+describe("AdditionalEnvVars", () => {
+  it("should parse single key-value pair", () => {
+    expect(AdditionalEnvVars.parse("FOO=bar")).toEqual({ FOO: "bar" });
+  });
+
+  it("should parse multiple key-value pairs", () => {
+    expect(AdditionalEnvVars.parse("FOO=bar,BAZ=qux")).toEqual({
+      FOO: "bar",
+      BAZ: "qux",
+    });
+  });
+
+  it("should handle whitespace", () => {
+    expect(AdditionalEnvVars.parse(" FOO = bar , BAZ = qux ")).toEqual({
+      FOO: "bar",
+      BAZ: "qux",
+    });
+  });
+
+  it("should return undefined for empty string", () => {
+    expect(AdditionalEnvVars.parse("")).toBeUndefined();
+  });
+
+  it("should return undefined for invalid format", () => {
+    expect(AdditionalEnvVars.parse("invalid")).toBeUndefined();
+  });
+
+  it("should skip invalid pairs but include valid ones", () => {
+    expect(AdditionalEnvVars.parse("FOO=bar,INVALID,BAZ=qux")).toEqual({
+      FOO: "bar",
+      BAZ: "qux",
+    });
+  });
+
+  it("should pass through undefined", () => {
+    expect(AdditionalEnvVars.parse(undefined)).toBeUndefined();
+  });
+
+  it("should handle empty values", () => {
+    expect(AdditionalEnvVars.parse("FOO=,BAR=value")).toEqual({
+      BAR: "value",
+    });
+  });
+});

apps/supervisor/src/envUtil.ts

@@ -0,0 +1,39 @@
+import { z } from "zod";
+
+export const BoolEnv = z.preprocess((val) => {
+  if (typeof val !== "string") {
+    return val;
+  }
+
+  return ["true", "1"].includes(val.toLowerCase().trim());
+}, z.boolean());
+
+export const AdditionalEnvVars = z.preprocess((val) => {
+  if (typeof val !== "string") {
+    return val;
+  }
+
+  if (!val) {
+    return undefined;
+  }
+
+  try {
+    const result = val.split(",").reduce(
+      (acc, pair) => {
+        const [key, value] = pair.split("=");
+        if (!key || !value) {
+          return acc;
+        }
+        acc[key.trim()] = value.trim();
+        return acc;
+      },
+      {} as Record<string, string>
+    );
+
+    // Return undefined if no valid key-value pairs were found
+    return Object.keys(result).length === 0 ? undefined : result;
+  } catch (error) {
+    console.warn("Failed to parse additional env vars", { error, val });
+    return undefined;
+  }
+}, z.record(z.string(), z.string()).optional());

apps/supervisor/src/index.ts

@@ -64,6 +64,7 @@ class ManagedSupervisor {
       imagePullSecrets: env.KUBERNETES_IMAGE_PULL_SECRETS?.split(","),
       heartbeatIntervalSeconds: env.RUNNER_HEARTBEAT_INTERVAL_SECONDS,
       snapshotPollIntervalSeconds: env.RUNNER_SNAPSHOT_POLL_INTERVAL_SECONDS,
+      additionalEnvVars: env.RUNNER_ADDITIONAL_ENV_VARS,
     } satisfies WorkloadManagerOptions;
 
     if (this.isKubernetes) {

apps/supervisor/src/workloadManager/docker.ts

@@ -57,6 +57,12 @@ export class DockerWorkloadManager implements WorkloadManager {
       );
     }
 
+    if (this.opts.additionalEnvVars) {
+      Object.entries(this.opts.additionalEnvVars).forEach(([key, value]) => {
+        runArgs.push(`--env=${key}=${value}`);
+      });
+    }
+
     if (env.ENFORCE_MACHINE_PRESETS) {
       runArgs.push(`--cpus=${opts.machine.cpu}`, `--memory=${opts.machine.memory}G`);
       runArgs.push(`--env=TRIGGER_MACHINE_CPU=${opts.machine.cpu}`);

apps/supervisor/src/workloadManager/kubernetes.ts

@@ -150,6 +150,12 @@ export class KubernetesWorkloadManager implements WorkloadManager {
                         },
                       ]
                     : []),
+                  ...(this.opts.additionalEnvVars
+                    ? Object.entries(this.opts.additionalEnvVars).map(([key, value]) => ({
+                        name: key,
+                        value: value,
+                      }))
+                    : []),
                 ],
               },
             ],

apps/supervisor/src/workloadManager/types.ts

@@ -8,6 +8,7 @@ export interface WorkloadManagerOptions {
   imagePullSecrets?: string[];
   heartbeatIntervalSeconds?: number;
   snapshotPollIntervalSeconds?: number;
+  additionalEnvVars?: Record<string, string>;
 }
 
 export interface WorkloadManager {