feat: add usage tracking to errored jobs

src/lib/mux-ai-error.ts

@@ -1,4 +1,7 @@
+import type { TokenUsage } from "../types.ts";
+
 import { detectLeakReason } from "./output-safety.ts";
+import { getErrorTokenUsage } from "./token-usage.ts";
 
 export type MuxAiErrorType = "validation_error" | "processing_error" | "timeout_error";
 
@@ -79,5 +82,11 @@ export function wrapError(error: unknown, message: string): never {
   if (shouldSuppress) {
     console.warn(`[@mux/ai] Suppressed suspected prompt leak in wrapped error (context: ${message}, reason: ${leakReason}).`);
   }
-  throw new Error(`${message}: ${detail}`);
+  const wrapped = new Error(`${message}: ${detail}`);
+  // Carry `usage` through wrapping so failed workflows still report tokens burned.
+  const usage = getErrorTokenUsage(error);
+  if (usage) {
+    (wrapped as Error & { usage?: TokenUsage }).usage = usage;
+  }
+  throw wrapped;
 }

src/lib/token-usage.ts

@@ -0,0 +1,91 @@
+import type { TokenUsage } from "../types.ts";
+
+const TOKEN_USAGE_FIELDS = [
+  "inputTokens",
+  "outputTokens",
+  "totalTokens",
+  "reasoningTokens",
+  "cachedInputTokens",
+] as const;
+
+type AggregatedTokenUsageField = (typeof TOKEN_USAGE_FIELDS)[number];
+
+function isDefinedTokenUsageValue(value: number | undefined): value is number {
+  return typeof value === "number";
+}
+
+export function aggregateTokenUsage(usages: TokenUsage[]): TokenUsage {
+  return TOKEN_USAGE_FIELDS.reduce<TokenUsage>((aggregate, field) => {
+    // Only aggregate values that were explicitly reported by the provider so
+    // omitted fields stay undefined instead of being coerced to 0.
+    const values = usages
+      .map(usage => usage[field as AggregatedTokenUsageField])
+      .filter(isDefinedTokenUsageValue);
+
+    if (values.length > 0) {
+      // Sum this field independently and write it back only when at least one
+      // chunk included real data for it.
+      aggregate[field] = values.reduce((total, value) => total + value, 0);
+    }
+
+    return aggregate;
+  }, {});
+}
+
+/**
+ * Reads token usage from a thrown error's plain `usage` property (AI SDK
+ * errors like `NoObjectGeneratedError`, or errors annotated by
+ * {@link rethrowWithTokenUsage}). Falls back to AI SDK v6's nested
+ * `inputTokenDetails`/`outputTokenDetails` for cache/reasoning counts.
+ */
+export function getErrorTokenUsage(error: unknown): TokenUsage | undefined {
+  if (typeof error !== "object" || error === null) {
+    return undefined;
+  }
+  const usage = (error as { usage?: unknown }).usage;
+  if (typeof usage !== "object" || usage === null) {
+    return undefined;
+  }
+
+  const source = usage as Record<string, unknown>;
+  const extracted: TokenUsage = {};
+  for (const field of TOKEN_USAGE_FIELDS) {
+    const value = source[field];
+    if (typeof value === "number") {
+      extracted[field] = value;
+    }
+  }
+
+  if (extracted.reasoningTokens === undefined) {
+    const reasoningTokens = (source.outputTokenDetails as { reasoningTokens?: unknown } | undefined)?.reasoningTokens;
+    if (typeof reasoningTokens === "number") {
+      extracted.reasoningTokens = reasoningTokens;
+    }
+  }
+  if (extracted.cachedInputTokens === undefined) {
+    const cacheReadTokens = (source.inputTokenDetails as { cacheReadTokens?: unknown } | undefined)?.cacheReadTokens;
+    if (typeof cacheReadTokens === "number") {
+      extracted.cachedInputTokens = cacheReadTokens;
+    }
+  }
+
+  return Object.keys(extracted).length > 0 ? extracted : undefined;
+}
+
+/**
+ * Rethrows `error` with the aggregate of `collectedUsage` plus the error's
+ * own usage attached as a plain enumerable `usage` property (mirrors the AI
+ * SDK's `NoObjectGeneratedError.usage` convention). Rethrows untouched when
+ * there is no usage to attach.
+ */
+export function rethrowWithTokenUsage(error: unknown, collectedUsage: TokenUsage[]): never {
+  const usages = [...collectedUsage];
+  const errorUsage = getErrorTokenUsage(error);
+  if (errorUsage) {
+    usages.push(errorUsage);
+  }
+  if (usages.length > 0 && typeof error === "object" && error !== null) {
+    (error as { usage?: TokenUsage }).usage = aggregateTokenUsage(usages);
+  }
+  throw error;
+}

src/workflows/ask-questions.ts

@@ -25,6 +25,7 @@ import {
 import { createLanguageModelFromConfig, resolveLanguageModelConfig } from "../lib/providers.ts";
 import type { ModelIdByProvider, SupportedProvider } from "../lib/providers.ts";
 import { withRetry } from "../lib/retry.ts";
+import { rethrowWithTokenUsage } from "../lib/token-usage.ts";
 import { resolveMuxSigningContext } from "../lib/workflow-credentials.ts";
 import { getStoryboardUrl } from "../primitives/storyboards.ts";
 import { fetchTranscriptForAsset } from "../primitives/transcripts.ts";
@@ -754,7 +755,20 @@ export async function askQuestions(
   options?: AskQuestionsOptions,
 ): Promise<AskQuestionsResult> {
   "use workflow";
+  const collectedUsage: TokenUsage[] = [];
+  try {
+    return await askQuestionsInternal(assetId, questions, options, collectedUsage);
+  } catch (error) {
+    rethrowWithTokenUsage(error, collectedUsage);
+  }
+}
 
+async function askQuestionsInternal(
+  assetId: string,
+  questions: Question[],
+  options: AskQuestionsOptions | undefined,
+  collectedUsage: TokenUsage[],
+): Promise<AskQuestionsResult> {
   // Validate questions array is non-empty
   if (!questions || questions.length === 0) {
     throw new MuxAiError("At least one question must be provided.", { type: "validation_error" });
@@ -955,6 +969,8 @@ export async function askQuestions(
     wrapError(error, `Failed to analyze ${contentType} questions with ${provider}`);
   }
 
+  collectedUsage.push(analysisResponse.usage);
+
   if (!analysisResponse.result?.answers) {
     throw new MuxAiError(`Failed to generate answers for asset ${assetId}.`);
   }

src/workflows/burned-in-captions.ts

@@ -21,6 +21,7 @@ import {
 } from "../lib/prompt-fragments.ts";
 import { createLanguageModelFromConfig, resolveLanguageModelConfig } from "../lib/providers.ts";
 import type { ModelIdByProvider, SupportedProvider } from "../lib/providers.ts";
+import { rethrowWithTokenUsage } from "../lib/token-usage.ts";
 import { getStoryboardUrl } from "../primitives/storyboards.ts";
 import type {
   ImageSubmissionMode,
@@ -328,6 +329,19 @@ export async function hasBurnedInCaptions(
   options: BurnedInCaptionsOptions = {},
 ): Promise<BurnedInCaptionsResult> {
   "use workflow";
+  const collectedUsage: TokenUsage[] = [];
+  try {
+    return await hasBurnedInCaptionsInternal(assetId, options, collectedUsage);
+  } catch (error) {
+    rethrowWithTokenUsage(error, collectedUsage);
+  }
+}
+
+async function hasBurnedInCaptionsInternal(
+  assetId: string,
+  options: BurnedInCaptionsOptions,
+  collectedUsage: TokenUsage[],
+): Promise<BurnedInCaptionsResult> {
   const {
     provider = DEFAULT_PROVIDER,
     model,
@@ -374,6 +388,8 @@ export async function hasBurnedInCaptions(
     });
   }
 
+  collectedUsage.push(analysisResponse.usage);
+
   if (!analysisResponse.result) {
     throw new Error("No analysis result received from AI provider");
   }

src/workflows/chapters.ts

@@ -21,6 +21,7 @@ import {
 import { createLanguageModelFromConfig, resolveLanguageModelConfig } from "../lib/providers.ts";
 import type { ModelIdByProvider, SupportedProvider } from "../lib/providers.ts";
 import { withRetry } from "../lib/retry.ts";
+import { rethrowWithTokenUsage } from "../lib/token-usage.ts";
 import { resolveMuxSigningContext } from "../lib/workflow-credentials.ts";
 import {
   extractTimestampedTranscript,
@@ -368,6 +369,19 @@ export async function generateChapters(
   options: ChaptersOptions = {},
 ): Promise<ChaptersResult> {
   "use workflow";
+  const collectedUsage: TokenUsage[] = [];
+  try {
+    return await generateChaptersInternal(assetId, options, collectedUsage);
+  } catch (error) {
+    rethrowWithTokenUsage(error, collectedUsage);
+  }
+}
+
+async function generateChaptersInternal(
+  assetId: string,
+  options: ChaptersOptions,
+  collectedUsage: TokenUsage[],
+): Promise<ChaptersResult> {
   const {
     languageCode,
     provider = "openai",
@@ -456,6 +470,10 @@ export async function generateChapters(
     wrapError(error, `Failed to generate chapters with ${provider}`);
   }
 
+  if (chaptersData) {
+    collectedUsage.push(chaptersData.usage);
+  }
+
   if (!chaptersData || !chaptersData.chapters) {
     throw new MuxAiError(`Failed to generate chapters for asset ${assetId}.`);
   }

src/workflows/edit-captions.ts

@@ -23,6 +23,7 @@ import {
   createPresignedGetUrlWithStorageAdapter,
   putObjectWithStorageAdapter,
 } from "../lib/storage-adapter.ts";
+import { rethrowWithTokenUsage } from "../lib/token-usage.ts";
 import {
   resolveMuxClient,
   resolveMuxSigningContext,
@@ -550,7 +551,20 @@ export async function editCaptions<P extends SupportedProvider = SupportedProvid
   options: EditCaptionsOptions<P>,
 ): Promise<EditCaptionsResult> {
   "use workflow";
+  const collectedUsage: TokenUsage[] = [];
+  try {
+    return await editCaptionsInternal(assetId, trackId, options, collectedUsage);
+  } catch (error) {
+    rethrowWithTokenUsage(error, collectedUsage);
+  }
+}
 
+async function editCaptionsInternal<P extends SupportedProvider = SupportedProvider>(
+  assetId: string,
+  trackId: string,
+  options: EditCaptionsOptions<P>,
+  collectedUsage: TokenUsage[],
+): Promise<EditCaptionsResult> {
   const {
     provider,
     model,
@@ -667,6 +681,7 @@ export async function editCaptions<P extends SupportedProvider = SupportedProvid
       });
       detectedProfanity = result.profanity;
       usage = result.usage;
+      collectedUsage.push(result.usage);
       // Record schema-smuggling signals from the step. zod.strip() has
       // already removed extras from the parsed output; the safety report
       // surfaces what was stripped.

src/workflows/embeddings.ts

@@ -7,6 +7,7 @@ import {
 import type { EmbeddingModelIdByProvider, SupportedEmbeddingProvider } from "../lib/providers.ts";
 import { createEmbeddingModelFromConfig, resolveEmbeddingModelConfig } from "../lib/providers.ts";
 import { withRetry } from "../lib/retry.ts";
+import { getErrorTokenUsage, rethrowWithTokenUsage } from "../lib/token-usage.ts";
 import { resolveMuxSigningContext } from "../lib/workflow-credentials.ts";
 import { chunkText, chunkVTTCues } from "../primitives/text-chunking.ts";
 import { fetchTranscriptForAsset, parseVTTCues } from "../primitives/transcripts.ts";
@@ -15,6 +16,7 @@ import type {
   ChunkingStrategy,
   MuxAIOptions,
   TextChunk,
+  TokenUsage,
   VideoEmbeddingsResult,
   WorkflowCredentialsInput,
 } from "../types.ts";
@@ -91,7 +93,7 @@ async function generateSingleChunkEmbedding({
   provider: SupportedEmbeddingProvider;
   modelId: string;
   credentials?: WorkflowCredentialsInput;
-}): Promise<ChunkEmbedding> {
+}): Promise<{ chunk: ChunkEmbedding; usage?: TokenUsage }> {
   "use step";
 
   const model = await createEmbeddingModelFromConfig(provider, modelId, credentials);
@@ -103,13 +105,18 @@ async function generateSingleChunkEmbedding({
   );
 
   return {
-    chunkId: chunk.id,
-    embedding: response.embedding,
-    metadata: {
-      startTime: chunk.startTime,
-      endTime: chunk.endTime,
-      tokenCount: chunk.tokenCount,
+    chunk: {
+      chunkId: chunk.id,
+      embedding: response.embedding,
+      metadata: {
+        startTime: chunk.startTime,
+        endTime: chunk.endTime,
+        tokenCount: chunk.tokenCount,
+      },
     },
+    usage: typeof response.usage?.tokens === "number" ?
+        { inputTokens: response.usage.tokens, totalTokens: response.usage.tokens } :
+      undefined,
   };
 }
 
@@ -146,7 +153,8 @@ async function generateSingleChunkEmbedding({
  */
 async function generateEmbeddingsInternal(
   assetId: string,
-  options: EmbeddingsOptions = {},
+  options: EmbeddingsOptions,
+  collectedUsage: TokenUsage[],
 ): Promise<EmbeddingsResult> {
   const {
     provider = "openai",
@@ -201,7 +209,7 @@ async function generateEmbeddingsInternal(
     for (let i = 0; i < chunks.length; i += batchSize) {
       const batch = chunks.slice(i, i + batchSize);
 
-      const batchResults = await Promise.all(
+      const batchOutcomes = await Promise.allSettled(
         batch.map(chunk =>
           generateSingleChunkEmbedding({
             chunk,
@@ -212,7 +220,26 @@ async function generateEmbeddingsInternal(
         ),
       );
 
-      chunkEmbeddings.push(...batchResults);
+      // The catch below replaces the error with a message-only Error, so
+      // usage from every settled call must be recorded here to survive.
+      const failures: unknown[] = [];
+      for (const outcome of batchOutcomes) {
+        if (outcome.status === "fulfilled") {
+          if (outcome.value.usage) {
+            collectedUsage.push(outcome.value.usage);
+          }
+          chunkEmbeddings.push(outcome.value.chunk);
+        } else {
+          const failureUsage = getErrorTokenUsage(outcome.reason);
+          if (failureUsage) {
+            collectedUsage.push(failureUsage);
+          }
+          failures.push(outcome.reason);
+        }
+      }
+      if (failures.length > 0) {
+        throw failures[0];
+      }
     }
   } catch (error) {
     throw new Error(
@@ -256,7 +283,12 @@ export async function generateEmbeddings(
   options: EmbeddingsOptions = {},
 ): Promise<EmbeddingsResult> {
   "use workflow";
-  return generateEmbeddingsInternal(assetId, options);
+  const collectedUsage: TokenUsage[] = [];
+  try {
+    return await generateEmbeddingsInternal(assetId, options, collectedUsage);
+  } catch (error) {
+    rethrowWithTokenUsage(error, collectedUsage);
+  }
 }
 
 /**
@@ -268,5 +300,10 @@ export async function generateVideoEmbeddings(
 ): Promise<EmbeddingsResult> {
   "use workflow";
   console.warn("generateVideoEmbeddings is deprecated. Use generateEmbeddings instead.");
-  return generateEmbeddingsInternal(assetId, options);
+  const collectedUsage: TokenUsage[] = [];
+  try {
+    return await generateEmbeddingsInternal(assetId, options, collectedUsage);
+  } catch (error) {
+    rethrowWithTokenUsage(error, collectedUsage);
+  }
 }