fix(aws-lambda): account for TaskScheduled/TaskSucceeded in cost

jrusso1020 · claude · jrusso1020 · commit d8851cfbdb46 · 2026-05-21T22:28:36.000Z
The CDK construct compiles tasks.LambdaInvoke to the optimized
arn:aws:states:::lambda:invoke integration, which emits Task* history
events with the Lambda response wrapped in .Payload. getRenderProgress
was only listening for the older LambdaFunction* events, so every CDK-
deployed stack reported $0 total cost and zero invocations on success
— a high-visibility regression that only surfaced when we manually
walked SFN history during a cost-analysis sweep.

Add cases for TaskScheduled (count invocation), TaskSucceeded (parse
Payload + accumulate billed duration / frame counts), and TaskFailed
(record error). Keep the LambdaFunction* paths so anyone wiring the
raw lambda:invokeFunction.sync task type still works. Factor out the
shared FramesEncoded-attribution logic so both branches agree on the
"only RenderChunk frames count" rule.

Tests pin a real-shape regression: replay the inspector-launch
1080p/30fps history (1 Plan + 16 RenderChunks + 1 Assemble) and assert
lambdaUsd lands at ~$0.582 — matching the cost-analysis script's
direct read against SFN history.

Co-Authored-By: Claude Opus 4.7 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/packages/aws-lambda/src/sdk/getRenderProgress.test.ts b/packages/aws-lambda/src/sdk/getRenderProgress.test.ts
@@ -67,6 +67,60 @@ function stateExited(name: string, output?: unknown): HistoryEvent {
   } as HistoryEvent;
 }
 
+// The CDK construct compiles `tasks.LambdaInvoke` to the optimized
+// `lambda:invoke` integration, which emits Task* events with the Lambda
+// response wrapped in `.Payload`. These helpers fabricate the matching
+// shapes so the test suite covers the same event stream a real render
+// produces — without them, the cost/progress code regressed silently to
+// $0 against production stacks.
+function taskScheduled(): HistoryEvent {
+  return {
+    type: "TaskScheduled",
+    id: 1,
+    timestamp: new Date(),
+    taskScheduledEventDetails: {
+      resource: "invoke",
+      resourceType: "lambda",
+      region: "us-east-1",
+      parameters: "{}",
+    },
+  } as HistoryEvent;
+}
+
+function taskSucceeded(payload: unknown): HistoryEvent {
+  // Match the optimized integration's wire shape — handler payload at
+  // `.Payload`, plus the ExecutedVersion / StatusCode envelope SFN
+  // surfaces verbatim from Lambda.
+  return {
+    type: "TaskSucceeded",
+    id: 1,
+    timestamp: new Date(),
+    taskSucceededEventDetails: {
+      resource: "invoke",
+      resourceType: "lambda",
+      output: JSON.stringify({
+        ExecutedVersion: "$LATEST",
+        Payload: payload,
+        StatusCode: 200,
+      }),
+    },
+  } as HistoryEvent;
+}
+
+function taskFailed(error: string, cause: string): HistoryEvent {
+  return {
+    type: "TaskFailed",
+    id: 1,
+    timestamp: new Date(),
+    taskFailedEventDetails: {
+      resource: "invoke",
+      resourceType: "lambda",
+      error,
+      cause,
+    },
+  } as HistoryEvent;
+}
+
 describe("getRenderProgress", () => {
   it("reports 0 progress before Plan completes", async () => {
     const sfn = new FakeSFN();
@@ -222,6 +276,130 @@ describe("getRenderProgress", () => {
   it("requires executionArn", async () => {
     await expect(getRenderProgress({ executionArn: "" })).rejects.toThrow(/executionArn/);
   });
+
+  // Real-world production renders emit Task* events (optimized
+  // lambda:invoke integration), not LambdaFunction* events. Before this
+  // block the SDK counted zero invocations and reported $0 cost on
+  // every CDK-deployed stack — see the in-tree cost analysis doc that
+  // hand-walked SFN history to confirm the bug.
+  describe("optimized lambda:invoke integration", () => {
+    it("counts a single TaskSucceeded as one Lambda invocation", async () => {
+      const sfn = new FakeSFN();
+      sfn.historyPages = [
+        [
+          stateEntered("Plan"),
+          taskScheduled(),
+          taskSucceeded({ Action: "plan", TotalFrames: 240, DurationMs: 1_000 }),
+        ],
+      ];
+      const progress = await getRenderProgress({
+        executionArn: "arn",
+        defaultMemorySizeMb: 10_240,
+        sfn: sfn as unknown as SFNClient,
+      });
+      expect(progress.lambdasInvoked).toBe(1);
+      expect(progress.totalFrames).toBe(240);
+      // 1 invocation × 1s × 10GB × $0.0000166667/GB-s ≈ $0.000167; the
+      // cost code rounds to 4 decimals (≈ tenth-of-a-cent), so a precise
+      // assertion needs precision=4 not precision=6. Pre-fix this read 0.
+      expect(progress.costs.breakdown.lambdaUsd).toBeCloseTo(0.0002, 4);
+      expect(progress.costs.breakdown.lambdaUsd).toBeGreaterThan(0);
+    });
+
+    it("attributes RenderChunk FramesEncoded but ignores Plan/Assemble FramesEncoded", async () => {
+      const sfn = new FakeSFN();
+      sfn.historyPages = [
+        [
+          stateEntered("Plan"),
+          taskScheduled(),
+          taskSucceeded({ Action: "plan", TotalFrames: 100, DurationMs: 1_000 }),
+          stateEntered("RenderChunk"),
+          taskScheduled(),
+          taskSucceeded({ Action: "renderChunk", FramesEncoded: 50, DurationMs: 2_000 }),
+          stateEntered("Assemble"),
+          taskScheduled(),
+          taskSucceeded({
+            Action: "assemble",
+            FramesEncoded: 100, // would double-count if Assemble's count bled in
+            FileSize: 9_000_000,
+            OutputS3Uri: "s3://b/k.mp4",
+            DurationMs: 1_500,
+          }),
+        ],
+      ];
+      const progress = await getRenderProgress({
+        executionArn: "arn",
+        sfn: sfn as unknown as SFNClient,
+      });
+      expect(progress.framesRendered).toBe(50);
+      expect(progress.lambdasInvoked).toBe(3);
+    });
+
+    it("captures TaskFailed errors with the enclosing state name", async () => {
+      const sfn = new FakeSFN();
+      sfn.historyPages = [
+        [
+          stateEntered("RenderChunk"),
+          taskScheduled(),
+          taskFailed("Sandbox.Timedout", "Task timed out after 900.00 seconds"),
+        ],
+      ];
+      const progress = await getRenderProgress({
+        executionArn: "arn",
+        sfn: sfn as unknown as SFNClient,
+      });
+      expect(progress.errors).toEqual([
+        {
+          state: "RenderChunk",
+          error: "Sandbox.Timedout",
+          cause: "Task timed out after 900.00 seconds",
+        },
+      ]);
+    });
+
+    it("billed-seconds sum matches what the cost-analysis script computes", async () => {
+      // Real-shape regression: cost script summed `DurationMs` across plan
+      // + renderChunk + assemble TaskSucceeded events and pinned the
+      // inspector-launch 1080p/30fps run at ~$0.58. Replay the same
+      // shape and verify the SDK lands on the same number.
+      const sfn = new FakeSFN();
+      const renderChunkSucceeded = (frames: number, ms: number) => [
+        stateEntered("RenderChunk"),
+        taskScheduled(),
+        taskSucceeded({ Action: "renderChunk", FramesEncoded: frames, DurationMs: ms }),
+      ];
+      sfn.historyPages = [
+        [
+          stateEntered("Plan"),
+          taskScheduled(),
+          taskSucceeded({ Action: "plan", TotalFrames: 1349, DurationMs: 13_000 }),
+          // 16 chunks averaging ~217 s each; the actual sweep saw
+          // 3492.7 s billed total.
+          ...Array.from({ length: 16 }, () => renderChunkSucceeded(84, 217_000)).flat(),
+          stateEntered("Assemble"),
+          taskScheduled(),
+          taskSucceeded({
+            Action: "assemble",
+            FileSize: 81_000_000,
+            OutputS3Uri: "s3://b/k.mp4",
+            DurationMs: 7_700,
+          }),
+        ],
+      ];
+      sfn.describe.status = "SUCCEEDED";
+      const progress = await getRenderProgress({
+        executionArn: "arn",
+        defaultMemorySizeMb: 10_240,
+        sfn: sfn as unknown as SFNClient,
+      });
+      // 13s + 16×217s + 7.7s = 3492.7 s × 10GB × $0.0000166667/GB-s
+      // ≈ $0.582 — the inspector-launch 1080p/30fps real run cost was
+      // $0.5826 (Lambda only). The pre-fix SDK reported $0.
+      expect(progress.costs.breakdown.lambdaUsd).toBeCloseTo(0.582, 2);
+      expect(progress.framesRendered).toBe(84 * 16);
+      expect(progress.lambdasInvoked).toBe(18); // 1 plan + 16 chunks + 1 assemble
+    });
+  });
 });
 
 void DescribeExecutionCommand;
diff --git a/packages/aws-lambda/src/sdk/getRenderProgress.ts b/packages/aws-lambda/src/sdk/getRenderProgress.ts
@@ -73,7 +73,12 @@ export interface RenderProgress {
   framesRendered: number;
   /** `null` until Plan completes. */
   totalFrames: number | null;
-  /** Count of `LambdaFunctionScheduled` events seen in the history so far. */
+  /**
+   * Count of Lambda invocations scheduled so far. Counts both the optimized
+   * `lambda:invoke` task type's `TaskScheduled` events and the older raw
+   * `LambdaFunctionScheduled` events, so this is the total whichever
+   * integration the state machine uses.
+   */
   lambdasInvoked: number;
   costs: RenderCost;
   /** Final output object if Assemble succeeded; `null` otherwise. */
@@ -198,9 +203,45 @@ function summarizeHistory(events: HistoryEvent[], memoryMb: number): HistorySumm
         stateTransitions++;
         currentLambdaState = ev.stateEnteredEventDetails?.name ?? currentLambdaState;
         break;
+      // The CDK construct uses `tasks.LambdaInvoke`, which compiles to the
+      // optimized `arn:aws:states:::lambda:invoke` integration. That
+      // integration emits `TaskScheduled` / `TaskSucceeded` / `TaskFailed`
+      // events (with the Lambda response wrapped in `.Payload`), NOT the
+      // older `LambdaFunctionScheduled` / `LambdaFunctionSucceeded` /
+      // `LambdaFunctionFailed` events the raw `lambda:invokeFunction.sync`
+      // task type would emit. Both paths are handled so the SDK works
+      // against either flavor of the construct.
+      case "TaskScheduled":
+        if (ev.taskScheduledEventDetails?.resourceType === "lambda") {
+          lambdasInvoked++;
+        }
+        break;
       case "LambdaFunctionScheduled":
         lambdasInvoked++;
         break;
+      case "TaskSucceeded": {
+        if (ev.taskSucceededEventDetails?.resourceType !== "lambda") break;
+        // Optimized integration wraps the Lambda response: the actual
+        // handler payload is at `.Payload`. Anywhere we used to read
+        // `payload.X` on a LambdaFunctionSucceeded output, we now read
+        // `payload.Payload.X` here.
+        const wrapped = parseJson(ev.taskSucceededEventDetails?.output);
+        const payload = unwrapLambdaPayload(wrapped);
+        const billedDurationMs = inferBilledMs(payload);
+        lambdaInvocations.push({
+          billedDurationMs,
+          memorySizeMb: memoryMb,
+          estimated: billedDurationMs === 0,
+        });
+        applyPayloadFrameCounts(payload, currentLambdaState, (delta) => {
+          framesRendered += delta;
+        });
+        if (payload && typeof payload === "object") {
+          const obj = payload as Record<string, unknown>;
+          if (typeof obj.TotalFrames === "number") totalFrames = obj.TotalFrames;
+        }
+        break;
+      }
       case "LambdaFunctionSucceeded": {
         const payload = parseJson(ev.lambdaFunctionSucceededEventDetails?.output);
         const billedDurationMs = inferBilledMs(payload);
@@ -245,6 +286,14 @@ function summarizeHistory(events: HistoryEvent[], memoryMb: number): HistorySumm
           }
         }
         break;
+      case "TaskFailed":
+        if (ev.taskFailedEventDetails?.resourceType !== "lambda") break;
+        errors.push({
+          state: currentLambdaState ?? "<unknown>",
+          error: ev.taskFailedEventDetails?.error ?? "UNKNOWN",
+          cause: ev.taskFailedEventDetails?.cause ?? "",
+        });
+        break;
       case "LambdaFunctionFailed":
         errors.push({
           state: currentLambdaState ?? "<unknown>",
@@ -299,6 +348,45 @@ function parseJson(s: string | undefined): unknown {
   }
 }
 
+/**
+ * The optimized `lambda:invoke` task wraps the Lambda response so the
+ * `TaskSucceeded.output` JSON looks like
+ *
+ *     { ExecutedVersion, Payload: { …handler payload… }, StatusCode, … }
+ *
+ * The raw `lambda:invokeFunction.sync` integration (older style) puts the
+ * handler payload at the root. Return the inner `Payload` object when
+ * present, otherwise the input — so downstream extraction code reads
+ * `.TotalFrames`/`.FramesEncoded`/`.DurationMs` the same way regardless of
+ * which integration is wired up. A non-object input passes through
+ * unchanged so the type-narrowing checks below still work.
+ */
+function unwrapLambdaPayload(payload: unknown): unknown {
+  if (payload && typeof payload === "object" && "Payload" in payload) {
+    const inner = (payload as { Payload: unknown }).Payload;
+    if (inner && typeof inner === "object") return inner;
+  }
+  return payload;
+}
+
+/**
+ * Apply `FramesEncoded` from a Lambda success payload to the running
+ * counter, but only when the enclosing state is the chunk-render step.
+ * Plan and Assemble payloads also carry `FramesEncoded` — counting them
+ * would double-count once Assemble runs. Shared between the
+ * `TaskSucceeded` and `LambdaFunctionSucceeded` branches.
+ */
+function applyPayloadFrameCounts(
+  payload: unknown,
+  currentLambdaState: string | null,
+  bump: (delta: number) => void,
+): void {
+  if (currentLambdaState !== "RenderChunk") return;
+  if (!payload || typeof payload !== "object") return;
+  const obj = payload as Record<string, unknown>;
+  if (typeof obj.FramesEncoded === "number") bump(obj.FramesEncoded);
+}
+
 /**
  * Lambda success payloads from our handler include `DurationMs` — the
  * wall-clock the handler observed. We use it as a best-effort proxy
