update runner id on restore

Author: nicktrn

packages/cli-v3/src/entryPoints/managed-run-controller.ts

@@ -86,6 +86,7 @@ type Metadata = {
   TRIGGER_SNAPSHOT_POLL_INTERVAL_SECONDS: number | undefined;
   TRIGGER_SUCCESS_EXIT_CODE: number | undefined;
   TRIGGER_FAILURE_EXIT_CODE: number | undefined;
+  TRIGGER_RUNNER_ID: string | undefined;
 };
 
 class MetadataClient {
@@ -126,6 +127,8 @@ class ManagedRunController {
   private workerApiUrl: string;
   private workerInstanceName: string;
 
+  private runnerId: string;
+
   private successExitCode = env.TRIGGER_SUCCESS_EXIT_CODE;
   private failureExitCode = env.TRIGGER_FAILURE_EXIT_CODE;
 
@@ -144,6 +147,8 @@ class ManagedRunController {
 
     this.workerManifest = opts.workerManifest;
 
+    this.runnerId = env.TRIGGER_RUNNER_ID;
+
     this.heartbeatIntervalSeconds = env.TRIGGER_HEARTBEAT_INTERVAL_SECONDS;
     this.snapshotPollIntervalSeconds = env.TRIGGER_SNAPSHOT_POLL_INTERVAL_SECONDS;
 
@@ -156,7 +161,7 @@ class ManagedRunController {
 
     this.httpClient = new WorkloadHttpClient({
       workerApiUrl: this.workerApiUrl,
-      runnerId: env.TRIGGER_RUNNER_ID,
+      runnerId: this.runnerId,
       deploymentId: env.TRIGGER_DEPLOYMENT_ID,
       deploymentVersion: env.TRIGGER_DEPLOYMENT_VERSION,
       projectRef: env.TRIGGER_PROJECT_REF,
@@ -709,6 +714,11 @@ class ManagedRunController {
 
       this.httpClient.updateApiUrl(this.workerApiUrl);
     }
+
+    if (overrides.TRIGGER_RUNNER_ID) {
+      this.runnerId = overrides.TRIGGER_RUNNER_ID;
+      this.httpClient.updateRunnerId(this.runnerId);
+    }
   }
 
   private async startAndExecuteRunAttempt({

packages/core/src/v3/runEngineWorker/workload/http.ts

@@ -20,13 +20,13 @@ type WorkloadHttpClientOptions = WorkloadClientCommonOptions;
 
 export class WorkloadHttpClient {
   private apiUrl: string;
+  private runnerId: string;
   private readonly deploymentId: string;
-  private readonly defaultHeaders: Record<string, string>;
 
-  constructor(opts: WorkloadHttpClientOptions) {
+  constructor(private opts: WorkloadHttpClientOptions) {
     this.apiUrl = opts.workerApiUrl.replace(/\/$/, "");
-    this.defaultHeaders = getDefaultWorkloadHeaders(opts);
     this.deploymentId = opts.deploymentId;
+    this.runnerId = opts.runnerId;
 
     if (!this.apiUrl) {
       throw new Error("apiURL is required and needs to be a non-empty string");
@@ -41,14 +41,25 @@ export class WorkloadHttpClient {
     this.apiUrl = apiUrl.replace(/\/$/, "");
   }
 
+  updateRunnerId(runnerId: string) {
+    this.runnerId = runnerId;
+  }
+
+  defaultHeaders(): Record<string, string> {
+    return getDefaultWorkloadHeaders({
+      ...this.opts,
+      runnerId: this.runnerId,
+    });
+  }
+
   async heartbeatRun(runId: string, snapshotId: string, body?: WorkloadHeartbeatRequestBody) {
     return wrapZodFetch(
       WorkloadHeartbeatResponseBody,
       `${this.apiUrl}/api/v1/workload-actions/runs/${runId}/snapshots/${snapshotId}/heartbeat`,
       {
         method: "POST",
         headers: {
-          ...this.defaultHeaders,
+          ...this.defaultHeaders(),
           "Content-Type": "application/json",
         },
         body: JSON.stringify(body ?? {}),
@@ -63,7 +74,7 @@ export class WorkloadHttpClient {
       {
         method: "GET",
         headers: {
-          ...this.defaultHeaders,
+          ...this.defaultHeaders(),
         },
       }
     );
@@ -76,7 +87,7 @@ export class WorkloadHttpClient {
       {
         method: "GET",
         headers: {
-          ...this.defaultHeaders,
+          ...this.defaultHeaders(),
         },
       }
     );
@@ -93,7 +104,7 @@ export class WorkloadHttpClient {
       {
         method: "POST",
         headers: {
-          ...this.defaultHeaders,
+          ...this.defaultHeaders(),
         },
         body: JSON.stringify(body),
       }
@@ -111,7 +122,7 @@ export class WorkloadHttpClient {
       {
         method: "POST",
         headers: {
-          ...this.defaultHeaders,
+          ...this.defaultHeaders(),
         },
         body: JSON.stringify(body),
       }
@@ -125,7 +136,7 @@ export class WorkloadHttpClient {
       {
         method: "GET",
         headers: {
-          ...this.defaultHeaders,
+          ...this.defaultHeaders(),
         },
       }
     );
@@ -139,7 +150,7 @@ export class WorkloadHttpClient {
         {
           method: "POST",
           headers: {
-            ...this.defaultHeaders,
+            ...this.defaultHeaders(),
             "Content-Type": "application/json",
           },
           body: JSON.stringify(body),
@@ -161,7 +172,7 @@ export class WorkloadHttpClient {
       {
         method: "GET",
         headers: {
-          ...this.defaultHeaders,
+          ...this.defaultHeaders(),
         },
       }
     );