initialize usage manager after env overrides

Author: nicktrn

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

@@ -98,9 +98,6 @@ process.on("uncaughtException", function (error, origin) {
   }
 });
 
-const usageIntervalMs = getEnvVar("USAGE_HEARTBEAT_INTERVAL_MS");
-const usageEventUrl = getEnvVar("USAGE_EVENT_URL");
-const triggerJWT = getEnvVar("TRIGGER_JWT");
 const heartbeatIntervalMs = getEnvVar("HEARTBEAT_INTERVAL_MS");
 
 const standardLocalsManager = new StandardLocalsManager();
@@ -112,17 +109,8 @@ lifecycleHooks.setGlobalLifecycleHooksManager(standardLifecycleHooksManager);
 const standardRunTimelineMetricsManager = new StandardRunTimelineMetricsManager();
 runTimelineMetrics.setGlobalManager(standardRunTimelineMetricsManager);
 
-const devUsageManager = new DevUsageManager();
-const prodUsageManager = new ProdUsageManager(devUsageManager, {
-  heartbeatIntervalMs: usageIntervalMs ? parseInt(usageIntervalMs, 10) : undefined,
-  url: usageEventUrl,
-  jwt: triggerJWT,
-});
-
-usage.setGlobalUsageManager(prodUsageManager);
-timeout.setGlobalManager(new UsageTimeoutManager(devUsageManager));
-
 resourceCatalog.setGlobalResourceCatalog(new StandardResourceCatalog());
+
 const durableClock = new DurableClock();
 clock.setGlobalClock(durableClock);
 const runMetadataManager = new StandardMetadataManager(
@@ -258,6 +246,12 @@ const zodIpc = new ZodIpcConnection({
         });
       }
 
+      initializeUsageManager({
+        usageIntervalMs: getEnvVar("USAGE_HEARTBEAT_INTERVAL_MS"),
+        usageEventUrl: getEnvVar("USAGE_EVENT_URL"),
+        triggerJWT: getEnvVar("TRIGGER_JWT"),
+      });
+
       standardRunTimelineMetricsManager.registerMetricsFromExecution(metrics);
 
       console.log(`[${new Date().toISOString()}] Received EXECUTE_TASK_RUN`, execution);
@@ -509,7 +503,7 @@ async function flushAll(timeoutInMs: number = 10_000) {
 async function flushUsage(timeoutInMs: number = 10_000) {
   const now = performance.now();
 
-  await Promise.race([prodUsageManager.flush(), setTimeout(timeoutInMs)]);
+  await Promise.race([usage.flush(), setTimeout(timeoutInMs)]);
 
   const duration = performance.now() - now;
 
@@ -551,6 +545,26 @@ async function flushMetadata(timeoutInMs: number = 10_000) {
   };
 }
 
+function initializeUsageManager({
+  usageIntervalMs,
+  usageEventUrl,
+  triggerJWT,
+}: {
+  usageIntervalMs?: string;
+  usageEventUrl?: string;
+  triggerJWT?: string;
+}) {
+  const devUsageManager = new DevUsageManager();
+  const prodUsageManager = new ProdUsageManager(devUsageManager, {
+    heartbeatIntervalMs: usageIntervalMs ? parseInt(usageIntervalMs, 10) : undefined,
+    url: usageEventUrl,
+    jwt: triggerJWT,
+  });
+
+  usage.setGlobalUsageManager(prodUsageManager);
+  timeout.setGlobalManager(new UsageTimeoutManager(devUsageManager));
+}
+
 const managedWorkerRuntime = new ManagedRuntimeManager(zodIpc, true);
 
 runtime.setGlobalRuntimeManager(managedWorkerRuntime);