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

jrusso1020 · claude · jrusso1020 · commit 36ddd487b1f5 · 2026-05-22T19:09:13.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,54 @@ function stateExited(name: string, output?: unknown): HistoryEvent {
   } as HistoryEvent;
 }
 
+// Optimized `lambda:invoke` integration's wire shape: Task* events with
+// the handler payload at `.Payload`. Helpers below fabricate these so
+// tests cover the same events a real CDK-deployed render produces.
+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 {
+  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 +270,116 @@ describe("getRenderProgress", () => {
   it("requires executionArn", async () => {
     await expect(getRenderProgress({ executionArn: "" })).rejects.toThrow(/executionArn/);
   });
+
+  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);
+      // computeRenderCost rounds to 4 decimals; precision=4 not 6.
+      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("sums billed seconds across plan + chunks + assemble", async () => {
+      const sfn = new FakeSFN();
+      const renderChunkSucceeded = (frames: number, ms: number) => [
+        stateEntered("RenderChunk"),
+        taskScheduled(),
+        taskSucceeded({ Action: "renderChunk", FramesEncoded: frames, DurationMs: ms }),
+      ];
+      // 1 plan + 16 chunks @ 217s + 1 assemble = 3492.7s billed.
+      sfn.historyPages = [
+        [
+          stateEntered("Plan"),
+          taskScheduled(),
+          taskSucceeded({ Action: "plan", TotalFrames: 1349, DurationMs: 13_000 }),
+          ...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,
+      });
+      // 3492.7s × 10GB × $0.0000166667/GB-s ≈ $0.582.
+      expect(progress.costs.breakdown.lambdaUsd).toBeCloseTo(0.582, 2);
+      expect(progress.framesRendered).toBe(84 * 16);
+      expect(progress.lambdasInvoked).toBe(18);
+    });
+  });
 });
 
 void DescribeExecutionCommand;
diff --git a/packages/aws-lambda/src/sdk/getRenderProgress.ts b/packages/aws-lambda/src/sdk/getRenderProgress.ts
@@ -73,7 +73,7 @@ export interface RenderProgress {
   framesRendered: number;
   /** `null` until Plan completes. */
   totalFrames: number | null;
-  /** Count of `LambdaFunctionScheduled` events seen in the history so far. */
+  /** Total Lambda invocations scheduled so far (both optimized + raw task integrations). */
   lambdasInvoked: number;
   costs: RenderCost;
   /** Final output object if Assemble succeeded; `null` otherwise. */
@@ -198,9 +198,35 @@ function summarizeHistory(events: HistoryEvent[], memoryMb: number): HistorySumm
         stateTransitions++;
         currentLambdaState = ev.stateEnteredEventDetails?.name ?? currentLambdaState;
         break;
+      // Optimized `lambda:invoke` task emits Task* events; raw
+      // `lambda:invokeFunction.sync` emits LambdaFunction*. Handle both.
+      case "TaskScheduled":
+        if (ev.taskScheduledEventDetails?.resourceType === "lambda") {
+          lambdasInvoked++;
+        }
+        break;
       case "LambdaFunctionScheduled":
         lambdasInvoked++;
         break;
+      case "TaskSucceeded": {
+        if (ev.taskSucceededEventDetails?.resourceType !== "lambda") break;
+        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);
@@ -209,20 +235,12 @@ function summarizeHistory(events: HistoryEvent[], memoryMb: number): HistorySumm
           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;
-          if (typeof obj.FramesEncoded === "number") {
-            // Plan and Assemble also return FramesEncoded; count framesRendered
-            // only inside the RenderChunk state so we don't double-count
-            // it on the Assemble pass. Keyed off the enclosing state name
-            // (set by the matching StateEntered) rather than the payload's
-            // `Action` field — `Action` is part of the Lambda event
-            // contract and not load-bearing for state-machine identity.
-            if (currentLambdaState === "RenderChunk") {
-              framesRendered += obj.FramesEncoded;
-            }
-          }
         }
         break;
       }
@@ -245,6 +263,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 +325,37 @@ function parseJson(s: string | undefined): unknown {
   }
 }
 
+/**
+ * Optimized `lambda:invoke` wraps the Lambda response as
+ * `{ ExecutedVersion, Payload: {…handler payload…}, StatusCode }`. Raw
+ * `lambda:invokeFunction.sync` puts the handler payload at the root.
+ * Return the inner `Payload` when present so callers read the same fields
+ * either way.
+ */
+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;
+}
+
+/**
+ * Bump `framesRendered` only inside the `RenderChunk` state. Plan and
+ * Assemble also report `FramesEncoded`, so a state-blind add would
+ * double-count once Assemble runs.
+ */
+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
