From e26dcb207960d64bef98cea7e6ac951bce85fcbe Mon Sep 17 00:00:00 2001
From: nicktrn <55853254+nicktrn@users.noreply.github.com>
Date: Tue, 6 May 2025 16:39:23 +0100
Subject: [PATCH] run and snapshot id can be overridden

---
 .../cli-v3/src/entryPoints/managed/env.ts     | 24 ++++++++++++-------
 .../src/entryPoints/managed/overrides.ts      |  2 ++
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/packages/cli-v3/src/entryPoints/managed/env.ts b/packages/cli-v3/src/entryPoints/managed/env.ts
index d3e381fa2..87b548335 100644
--- a/packages/cli-v3/src/entryPoints/managed/env.ts
+++ b/packages/cli-v3/src/entryPoints/managed/env.ts
@@ -22,8 +22,6 @@ const Env = z.object({
   // Set at runtime
   TRIGGER_WORKLOAD_CONTROLLER_ID: z.string().default(`controller_${randomUUID()}`),
   TRIGGER_ENV_ID: z.string(),
-  TRIGGER_RUN_ID: z.string().optional(), // This is only useful for cold starts
-  TRIGGER_SNAPSHOT_ID: z.string().optional(), // This is only useful for cold starts
   OTEL_EXPORTER_OTLP_ENDPOINT: z.string().url(),
   TRIGGER_WARM_START_URL: z.string().optional(),
   TRIGGER_WARM_START_CONNECTION_TIMEOUT_MS: z.coerce.number().default(30_000),
@@ -38,6 +36,8 @@ const Env = z.object({
   TRIGGER_DEQUEUED_AT_MS: DateEnv,
 
   // May be overridden
+  TRIGGER_RUN_ID: z.string().optional(), // This is set for cold starts and restores
+  TRIGGER_SNAPSHOT_ID: z.string().optional(), // This is set for cold starts and restores
   TRIGGER_SUPERVISOR_API_PROTOCOL: z.enum(["http", "https"]),
   TRIGGER_SUPERVISOR_API_DOMAIN: z.string(),
   TRIGGER_SUPERVISOR_API_PORT: z.coerce.number(),
@@ -94,12 +94,6 @@ export class RunnerEnv {
   get TRIGGER_ENV_ID() {
     return this.env.TRIGGER_ENV_ID;
   }
-  get TRIGGER_RUN_ID() {
-    return this.env.TRIGGER_RUN_ID;
-  }
-  get TRIGGER_SNAPSHOT_ID() {
-    return this.env.TRIGGER_SNAPSHOT_ID;
-  }
   get TRIGGER_WARM_START_URL() {
     return this.env.TRIGGER_WARM_START_URL;
   }
@@ -126,6 +120,12 @@ export class RunnerEnv {
   }
 
   // Overridable values
+  get TRIGGER_RUN_ID() {
+    return this.env.TRIGGER_RUN_ID;
+  }
+  get TRIGGER_SNAPSHOT_ID() {
+    return this.env.TRIGGER_SNAPSHOT_ID;
+  }
   get TRIGGER_SUCCESS_EXIT_CODE() {
     return this.env.TRIGGER_SUCCESS_EXIT_CODE;
   }
@@ -163,6 +163,14 @@ export class RunnerEnv {
 
   /** Overrides existing env vars with new values */
   override(overrides: Metadata) {
+    if (overrides.TRIGGER_RUN_ID) {
+      this.env.TRIGGER_RUN_ID = overrides.TRIGGER_RUN_ID;
+    }
+
+    if (overrides.TRIGGER_SNAPSHOT_ID) {
+      this.env.TRIGGER_SNAPSHOT_ID = overrides.TRIGGER_SNAPSHOT_ID;
+    }
+
     if (overrides.TRIGGER_SUCCESS_EXIT_CODE) {
       this.env.TRIGGER_SUCCESS_EXIT_CODE = overrides.TRIGGER_SUCCESS_EXIT_CODE;
     }
diff --git a/packages/cli-v3/src/entryPoints/managed/overrides.ts b/packages/cli-v3/src/entryPoints/managed/overrides.ts
index 872b5ad0b..862e6c6fb 100644
--- a/packages/cli-v3/src/entryPoints/managed/overrides.ts
+++ b/packages/cli-v3/src/entryPoints/managed/overrides.ts
@@ -1,4 +1,6 @@
 export type Metadata = {
+  TRIGGER_RUN_ID: string | undefined;
+  TRIGGER_SNAPSHOT_ID: string | undefined;
   TRIGGER_SUPERVISOR_API_PROTOCOL: string | undefined;
   TRIGGER_SUPERVISOR_API_DOMAIN: string | undefined;
   TRIGGER_SUPERVISOR_API_PORT: number | undefined;