Fix usage by still passing durationMs and costInCents to the execution, just not the run.ctx

Author: ericallam

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

@@ -505,6 +505,11 @@ const zodIpc = new ZodIpcConnection({
             getNumberEnvVar("TRIGGER_RUN_METADATA_FLUSH_INTERVAL", 1000)
           );
 
+          devUsageManager.setInitialState({
+            cpuTime: execution.run.durationMs ?? 0,
+            costInCents: execution.run.costInCents ?? 0,
+          });
+
           _executionMeasurement = usage.start();
 
           const timeoutController = timeout.abortAfterTimeout(execution.run.maxDuration);

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

@@ -328,12 +328,17 @@ const zodIpc = new ZodIpcConnection({
 
       resetExecutionEnvironment();
 
-      initializeUsageManager({
+      const prodManager = initializeUsageManager({
         usageIntervalMs: getEnvVar("USAGE_HEARTBEAT_INTERVAL_MS"),
         usageEventUrl: getEnvVar("USAGE_EVENT_URL"),
         triggerJWT: getEnvVar("TRIGGER_JWT"),
       });
 
+      prodManager.setInitialState({
+        cpuTime: execution.run.durationMs ?? 0,
+        costInCents: execution.run.costInCents ?? 0,
+      });
+
       standardRunTimelineMetricsManager.registerMetricsFromExecution(metrics, isWarmStart);
 
       console.log(`[${new Date().toISOString()}] Received EXECUTE_TASK_RUN`, execution);
@@ -690,6 +695,8 @@ function initializeUsageManager({
 
   usage.setGlobalUsageManager(prodUsageManager);
   timeout.setGlobalManager(new UsageTimeoutManager(devUsageManager));
+
+  return prodUsageManager;
 }
 
 _sharedWorkerRuntime = new SharedRuntimeManager(zodIpc, true);

packages/core/src/v3/schemas/common.ts

@@ -225,6 +225,10 @@ export const TaskRun = z.object({
 
   parentTaskRunId: z.string().optional(),
   rootTaskRunId: z.string().optional(),
+
+  // These are only used during execution, not in run.ctx
+  durationMs: z.number().optional(),
+  costInCents: z.number().optional(),
 });
 
 export type TaskRun = z.infer<typeof TaskRun>;
@@ -397,6 +401,8 @@ export const TaskRunContext = z.object({
     payload: true,
     payloadType: true,
     metadata: true,
+    durationMs: true,
+    costInCents: true,
   }),
   ...StaticTaskRunExecutionShape,
 });

packages/core/src/v3/usage/api.ts

@@ -1,7 +1,7 @@
 const API_NAME = "usage";
 
 import { getGlobal, registerGlobal, unregisterGlobal } from "../utils/globals.js";
-import type { UsageManager, UsageMeasurement, UsageSample } from "./types.js";
+import type { InitialUsageState, UsageManager, UsageMeasurement, UsageSample } from "./types.js";
 import { NoopUsageManager } from "./noopUsageManager.js";
 
 const NOOP_USAGE_MANAGER = new NoopUsageManager();
@@ -53,6 +53,10 @@ export class UsageAPI implements UsageManager {
     this.disable();
   }
 
+  public getInitialState(): InitialUsageState {
+    return this.#getUsageManager().getInitialState();
+  }
+
   #getUsageManager(): UsageManager {
     return getGlobal(API_NAME) ?? NOOP_USAGE_MANAGER;
   }

packages/core/src/v3/usage/devUsageManager.ts

@@ -1,4 +1,4 @@
-import { UsageManager, UsageMeasurement, UsageSample } from "./types.js";
+import { InitialUsageState, UsageManager, UsageMeasurement, UsageSample } from "./types.js";
 import { clock } from "../clock-api.js";
 import { ClockTime, calculateDurationInMs } from "../clock/clock.js";
 
@@ -45,6 +45,10 @@ export class DevUsageManager implements UsageManager {
   private _firstMeasurement?: DevUsageMeasurement;
   private _currentMeasurements: Map<string, DevUsageMeasurement> = new Map();
   private _pauses: Map<string, { start: ClockTime; end?: ClockTime }> = new Map();
+  private _initialState: InitialUsageState = {
+    cpuTime: 0,
+    costInCents: 0,
+  };
 
   disable(): void {}
 
@@ -54,6 +58,18 @@ export class DevUsageManager implements UsageManager {
     this._firstMeasurement = undefined;
     this._currentMeasurements.clear();
     this._pauses.clear();
+    this._initialState = {
+      cpuTime: 0,
+      costInCents: 0,
+    };
+  }
+
+  setInitialState(state: InitialUsageState) {
+    this._initialState = state;
+  }
+
+  getInitialState(): InitialUsageState {
+    return this._initialState;
   }
 
   sample(): UsageSample | undefined {

packages/core/src/v3/usage/noopUsageManager.ts

@@ -1,4 +1,4 @@
-import { UsageManager, UsageMeasurement, UsageSample } from "./types.js";
+import { InitialUsageState, UsageManager, UsageMeasurement, UsageSample } from "./types.js";
 
 export class NoopUsageManager implements UsageManager {
   disable(): void {
@@ -30,4 +30,11 @@ export class NoopUsageManager implements UsageManager {
   reset(): void {
     // Noop
   }
+
+  getInitialState(): InitialUsageState {
+    return {
+      cpuTime: 0,
+      costInCents: 0,
+    };
+  }
 }

packages/core/src/v3/usage/prodUsageManager.ts

@@ -1,5 +1,5 @@
 import { setInterval } from "node:timers/promises";
-import { UsageManager, UsageMeasurement, UsageSample } from "./types.js";
+import { InitialUsageState, UsageManager, UsageMeasurement, UsageSample } from "./types.js";
 import { UsageClient } from "./usageClient.js";
 
 export type ProdUsageManagerOptions = {
@@ -13,6 +13,10 @@ export class ProdUsageManager implements UsageManager {
   private _abortController: AbortController | undefined;
   private _lastSample: UsageSample | undefined;
   private _usageClient: UsageClient | undefined;
+  private _initialState: InitialUsageState = {
+    cpuTime: 0,
+    costInCents: 0,
+  };
 
   constructor(
     private readonly delegageUsageManager: UsageManager,
@@ -27,13 +31,25 @@ export class ProdUsageManager implements UsageManager {
     return typeof this._usageClient !== "undefined";
   }
 
+  setInitialState(state: InitialUsageState) {
+    this._initialState = state;
+  }
+
+  getInitialState(): InitialUsageState {
+    return this._initialState;
+  }
+
   reset(): void {
     this.delegageUsageManager.reset();
     this._abortController?.abort();
     this._abortController = new AbortController();
     this._usageClient = undefined;
     this._measurement = undefined;
     this._lastSample = undefined;
+    this._initialState = {
+      cpuTime: 0,
+      costInCents: 0,
+    };
   }
 
   disable(): void {

packages/core/src/v3/usage/types.ts

@@ -7,8 +7,14 @@ export interface UsageMeasurement {
   sample(): UsageSample;
 }
 
+export type InitialUsageState = {
+  cpuTime: number;
+  costInCents: number;
+};
+
 export interface UsageManager {
   disable(): void;
+  getInitialState(): InitialUsageState;
   start(): UsageMeasurement;
   stop(measurement: UsageMeasurement): UsageSample;
   sample(): UsageSample | undefined;

packages/trigger-sdk/src/v3/usage.ts

@@ -54,6 +54,7 @@ export const usage = {
    */
   getCurrent: (): CurrentUsage => {
     const sample = usageApi.sample();
+    const initialState = usageApi.getInitialState();
     const machine = taskContext.ctx?.machine;
     const run = taskContext.ctx?.run;
 
@@ -65,12 +66,12 @@ export const usage = {
             durationMs: 0,
           },
           total: {
-            costInCents: 0,
-            durationMs: 0,
+            costInCents: initialState.costInCents,
+            durationMs: initialState.cpuTime,
           },
         },
         baseCostInCents: run?.baseCostInCents ?? 0,
-        totalCostInCents: run?.baseCostInCents ?? 0,
+        totalCostInCents: initialState.costInCents + (run?.baseCostInCents ?? 0),
       };
     }
 
@@ -83,12 +84,12 @@ export const usage = {
           durationMs: sample.cpuTime,
         },
         total: {
-          costInCents: currentCostInCents,
-          durationMs: sample.cpuTime,
+          costInCents: currentCostInCents + initialState.costInCents,
+          durationMs: sample.cpuTime + initialState.cpuTime,
         },
       },
       baseCostInCents: run?.baseCostInCents ?? 0,
-      totalCostInCents: currentCostInCents + (run?.baseCostInCents ?? 0),
+      totalCostInCents: currentCostInCents + (run?.baseCostInCents ?? 0) + initialState.costInCents,
     };
   },
   /**

references/hello-world/src/trigger/usage.ts

@@ -0,0 +1,35 @@
+import { logger, task, wait, usage } from "@trigger.dev/sdk";
+import { setTimeout } from "timers/promises";
+
+export const usageExampleTask = task({
+  id: "usage-example",
+  retry: {
+    maxAttempts: 3,
+    minTimeoutInMs: 500,
+    maxTimeoutInMs: 1000,
+    factor: 1.5,
+  },
+  run: async (payload: { throwError: boolean }, { ctx }) => {
+    logger.info("run.ctx", { ctx });
+
+    await setTimeout(1000);
+
+    const currentUsage = usage.getCurrent();
+
+    logger.info("currentUsage", { currentUsage });
+
+    if (payload.throwError && ctx.attempt.number === 1) {
+      throw new Error("Forced error to cause a retry");
+    }
+
+    await setTimeout(5000);
+
+    const currentUsage2 = usage.getCurrent();
+
+    logger.info("currentUsage2", { currentUsage2 });
+
+    return {
+      message: "Hello, world!",
+    };
+  },
+});