From b781213ac2e706328a2ac196a7b64dcef334cfc8 Mon Sep 17 00:00:00 2001 From: Walker Date: Mon, 13 Jul 2026 12:28:55 -0700 Subject: [PATCH 1/2] feat: add usage tracking to errored jobs --- src/lib/mux-ai-error.ts | 14 +- src/lib/token-usage.ts | 100 ++++++++++++ src/workflows/ask-questions.ts | 19 +++ src/workflows/burned-in-captions.ts | 18 +++ src/workflows/chapters.ts | 21 +++ src/workflows/edit-captions.ts | 18 +++ src/workflows/embeddings.ts | 61 +++++-- src/workflows/engagement-insights.ts | 17 ++ src/workflows/summarization.ts | 19 +++ src/workflows/translate-captions.ts | 145 ++++++++++------- tests/unit/token-usage.test.ts | 153 ++++++++++++++++++ .../translate-captions-error-usage.test.ts | 147 +++++++++++++++++ 12 files changed, 664 insertions(+), 68 deletions(-) create mode 100644 src/lib/token-usage.ts create mode 100644 tests/unit/token-usage.test.ts create mode 100644 tests/unit/translate-captions-error-usage.test.ts diff --git a/src/lib/mux-ai-error.ts b/src/lib/mux-ai-error.ts index d797678..5d4cb68 100644 --- a/src/lib/mux-ai-error.ts +++ b/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,14 @@ 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}`); + // Token usage rides on errors as a plain `usage` property (AI SDK errors + // like NoObjectGeneratedError, plus errors annotated by + // rethrowWithTokenUsage). Carry it through wrapping so failed workflows + // can still report the tokens they burned. + const usage = getErrorTokenUsage(error); + if (usage) { + (wrapped as Error & { usage?: TokenUsage }).usage = usage; + } + throw wrapped; } diff --git a/src/lib/token-usage.ts b/src/lib/token-usage.ts new file mode 100644 index 0000000..3129b4c --- /dev/null +++ b/src/lib/token-usage.ts @@ -0,0 +1,100 @@ +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((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 carried on a thrown error's plain `usage` property. + * + * Covers both AI SDK errors that report usage for the failed call (e.g. + * `NoObjectGeneratedError`) and errors this package has already annotated + * via {@link rethrowWithTokenUsage}. Only the known numeric fields are + * extracted; AI SDK v6 nests reasoning/cache counts under + * `outputTokenDetails`/`inputTokenDetails` when the deprecated flat fields + * are absent, so those are used as fallbacks. + */ +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; + 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 token usage of all provider calls made + * so far attached as a plain enumerable `usage` property, so consumers can + * report tokens burned by failed workflows (mirrors the AI SDK's + * `NoObjectGeneratedError.usage` convention). + * + * The error's own `usage` (from the failing call) is folded into the + * aggregate. When no usage was collected and the error carries none, the + * error is rethrown untouched — callers treat a missing `usage` as "no + * tokens spent". + */ +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; +} diff --git a/src/workflows/ask-questions.ts b/src/workflows/ask-questions.ts index 6681dcf..2da508c 100644 --- a/src/workflows/ask-questions.ts +++ b/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,23 @@ export async function askQuestions( options?: AskQuestionsOptions, ): Promise { "use workflow"; + // Usage from provider calls made so far. A throw after the analysis call + // (e.g. incomplete answers) still reports the tokens burned via the + // error's `usage` property. + 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 { // 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 +972,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}.`); } diff --git a/src/workflows/burned-in-captions.ts b/src/workflows/burned-in-captions.ts index de24e88..c1f12fd 100644 --- a/src/workflows/burned-in-captions.ts +++ b/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,21 @@ export async function hasBurnedInCaptions( options: BurnedInCaptionsOptions = {}, ): Promise { "use workflow"; + // Usage from provider calls made so far. A throw after the analysis call + // still reports the tokens burned via the error's `usage` property. + 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 { const { provider = DEFAULT_PROVIDER, model, @@ -374,6 +390,8 @@ export async function hasBurnedInCaptions( }); } + collectedUsage.push(analysisResponse.usage); + if (!analysisResponse.result) { throw new Error("No analysis result received from AI provider"); } diff --git a/src/workflows/chapters.ts b/src/workflows/chapters.ts index 0448ede..f5adc12 100644 --- a/src/workflows/chapters.ts +++ b/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,22 @@ export async function generateChapters( options: ChaptersOptions = {}, ): Promise { "use workflow"; + // Usage from provider calls made so far. A throw after the generation call + // (e.g. no valid chapters) still reports the tokens burned via the error's + // `usage` property. + 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 { const { languageCode, provider = "openai", @@ -456,6 +473,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}.`); } diff --git a/src/workflows/edit-captions.ts b/src/workflows/edit-captions.ts index 2b0fde4..5349353 100644 --- a/src/workflows/edit-captions.ts +++ b/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,23 @@ export async function editCaptions

, ): Promise { "use workflow"; + // Usage from provider calls made so far. A throw after profanity detection + // (e.g. a failed S3 upload) still reports the tokens burned via the + // error's `usage` property. + const collectedUsage: TokenUsage[] = []; + try { + return await editCaptionsInternal(assetId, trackId, options, collectedUsage); + } catch (error) { + rethrowWithTokenUsage(error, collectedUsage); + } +} +async function editCaptionsInternal

( + assetId: string, + trackId: string, + options: EditCaptionsOptions

, + collectedUsage: TokenUsage[], +): Promise { const { provider, model, @@ -667,6 +684,7 @@ export async function editCaptions

{ +}): Promise<{ chunk: ChunkEmbedding; usage?: TokenUsage }> { "use step"; const model = await createEmbeddingModelFromConfig(provider, modelId, credentials); @@ -103,13 +105,21 @@ 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, + }, }, + // Embedding tokens are all input tokens. Reported so a failure partway + // through the batches can attach the tokens burned so far to the error; + // the success-path result intentionally remains unchanged. + usage: typeof response.usage?.tokens === "number" ? + { inputTokens: response.usage.tokens, totalTokens: response.usage.tokens } : + undefined, }; } @@ -146,7 +156,8 @@ async function generateSingleChunkEmbedding({ */ async function generateEmbeddingsInternal( assetId: string, - options: EmbeddingsOptions = {}, + options: EmbeddingsOptions, + collectedUsage: TokenUsage[], ): Promise { const { provider = "openai", @@ -201,7 +212,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 +223,23 @@ async function generateEmbeddingsInternal( ), ); - chunkEmbeddings.push(...batchResults); + // Record usage from every settled call — including the fulfilled + // siblings of a failed call — before rethrowing, so the tokens + // burned so far ride on the thrown error. + 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 { + 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 { "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 { "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); + } } diff --git a/src/workflows/engagement-insights.ts b/src/workflows/engagement-insights.ts index c18f0fe..3e37cf4 100644 --- a/src/workflows/engagement-insights.ts +++ b/src/workflows/engagement-insights.ts @@ -26,6 +26,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 { signUrl } from "../lib/url-signing.ts"; import { resolveMuxSigningContext } from "../lib/workflow-credentials.ts"; import type { HeatmapResponse } from "../primitives/heatmap.ts"; @@ -642,7 +643,22 @@ export async function generateEngagementInsights( options: EngagementInsightsOptions = {}, ): Promise { "use workflow"; + // Usage from provider calls made so far. A throw after the insights call + // (e.g. no valid insights) still reports the tokens burned via the error's + // `usage` property. + const collectedUsage: TokenUsage[] = []; + try { + return await generateEngagementInsightsInternal(assetId, options, collectedUsage); + } catch (error) { + rethrowWithTokenUsage(error, collectedUsage); + } +} +async function generateEngagementInsightsInternal( + assetId: string, + options: EngagementInsightsOptions, + collectedUsage: TokenUsage[], +): Promise { const { provider = "openai", model, @@ -819,6 +835,7 @@ export async function generateEngagementInsights( credentials, ); const { result: aiInsights, usage } = aiStepResult; + collectedUsage.push(usage); if (!aiInsights.momentInsights || aiInsights.momentInsights.length === 0) { throw new MuxAiError(`Failed to generate insights for asset ${assetId}.`); diff --git a/src/workflows/summarization.ts b/src/workflows/summarization.ts index c3ba007..0b6ebf2 100644 --- a/src/workflows/summarization.ts +++ b/src/workflows/summarization.ts @@ -39,6 +39,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"; @@ -683,6 +684,22 @@ export async function getSummaryAndTags( options?: SummarizationOptions, ): Promise { "use workflow"; + // Usage from provider calls made so far. A throw after the analysis call + // (e.g. unusable model output) still reports the tokens burned via the + // error's `usage` property. + const collectedUsage: TokenUsage[] = []; + try { + return await getSummaryAndTagsInternal(assetId, options, collectedUsage); + } catch (error) { + rethrowWithTokenUsage(error, collectedUsage); + } +} + +async function getSummaryAndTagsInternal( + assetId: string, + options: SummarizationOptions | undefined, + collectedUsage: TokenUsage[], +): Promise { const { provider = "openai", model, @@ -830,6 +847,8 @@ export async function getSummaryAndTags( wrapError(error, `Failed to analyze ${contentType} content with ${provider}`); } + collectedUsage.push(analysisResponse.usage); + if (!analysisResponse.result) { const contentType = isAudioOnly ? "audio" : "video"; throw new MuxAiError(`Failed to analyze ${contentType} content for asset ${assetId}.`); diff --git a/src/workflows/translate-captions.ts b/src/workflows/translate-captions.ts index 654c18a..3120841 100644 --- a/src/workflows/translate-captions.ts +++ b/src/workflows/translate-captions.ts @@ -35,6 +35,7 @@ import { createPresignedGetUrlWithStorageAdapter, putObjectWithStorageAdapter, } from "../lib/storage-adapter.ts"; +import { aggregateTokenUsage, getErrorTokenUsage, rethrowWithTokenUsage } from "../lib/token-usage.ts"; import { resolveMuxSigningContext } from "../lib/workflow-credentials.ts"; import { chunkVTTCuesByBudget, @@ -233,15 +234,8 @@ interface TranslationChunkRequest { cueBlocks: string[]; } -const TOKEN_USAGE_FIELDS = [ - "inputTokens", - "outputTokens", - "totalTokens", - "reasoningTokens", - "cachedInputTokens", -] as const; - -type AggregatedTokenUsageField = (typeof TOKEN_USAGE_FIELDS)[number]; +// Re-exported from its original home for backwards compatibility. +export { aggregateTokenUsage }; class TranslationChunkValidationError extends Error { constructor(message: string) { @@ -286,10 +280,6 @@ export function shouldSplitChunkTranslationError(error: unknown): boolean { ); } -function isDefinedTokenUsageValue(value: number | undefined): value is number { - return typeof value === "number"; -} - function resolveTranslationChunkingOptions( options?: TranslationChunkingOptions, ): Required { @@ -320,22 +310,39 @@ function resolveTranslationChunkingOptions( }; } -export function aggregateTokenUsage(usages: TokenUsage[]): TokenUsage { - return TOKEN_USAGE_FIELDS.reduce((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); +/** + * Unwraps `Promise.allSettled` outcomes for a group of chunk translations. + * On any rejection, the usage of the fulfilled siblings (plus `priorUsage` + * and the usage carried by the other rejections) is attached to the first + * rejection before it is rethrown, so tokens burned on successful chunks + * aren't lost when one chunk fails. + */ +function collectSettledTranslationsOrRethrow( + outcomes: Array>, + priorUsage: TokenUsage[], +): TranslateStepResult[] { + const results: TranslateStepResult[] = []; + const failures: unknown[] = []; + for (const outcome of outcomes) { + if (outcome.status === "fulfilled") { + results.push(outcome.value); + } else { + failures.push(outcome.reason); + } + } + + if (failures.length > 0) { + const partialUsage = [...priorUsage, ...results.map(result => result.usage)]; + for (const failure of failures.slice(1)) { + const usage = getErrorTokenUsage(failure); + if (usage) { + partialUsage.push(usage); + } } + rethrowWithTokenUsage(failures[0], partialUsage); + } - return aggregate; - }, {}); + return results; } function createTranslationChunkRequest( @@ -838,24 +845,30 @@ async function translateChunkWithFallback({ } const [leftChunk, rightChunk] = splitTranslationChunkAtMidpoint(chunk); - const [leftResult, rightResult] = await Promise.all([ - translateChunkWithFallback({ - chunk: leftChunk, - fromLanguageCode, - toLanguageCode, - provider, - modelId, - credentials, - }), - translateChunkWithFallback({ - chunk: rightChunk, - fromLanguageCode, - toLanguageCode, - provider, - modelId, - credentials, - }), - ]); + // The failed whole-chunk attempt burned tokens too; fold its usage into + // the error if the split halves also fail. + const failedAttemptUsage = getErrorTokenUsage(error); + const [leftResult, rightResult] = collectSettledTranslationsOrRethrow( + await Promise.allSettled([ + translateChunkWithFallback({ + chunk: leftChunk, + fromLanguageCode, + toLanguageCode, + provider, + modelId, + credentials, + }), + translateChunkWithFallback({ + chunk: rightChunk, + fromLanguageCode, + toLanguageCode, + provider, + modelId, + credentials, + }), + ]), + failedAttemptUsage ? [failedAttemptUsage] : [], + ); return { translatedVtt: concatenateVttSegments([leftResult.translatedVtt, rightResult.translatedVtt]), @@ -910,17 +923,20 @@ async function translateCaptionTrack({ for (let index = 0; index < chunkPlan.chunks.length; index += resolvedChunking.maxConcurrentTranslations) { const batch = chunkPlan.chunks.slice(index, index + resolvedChunking.maxConcurrentTranslations); - const batchResults = await Promise.all( - batch.map(chunk => - translateChunkWithFallback({ - chunk, - fromLanguageCode, - toLanguageCode, - provider, - modelId, - credentials, - }), + const batchResults = collectSettledTranslationsOrRethrow( + await Promise.allSettled( + batch.map(chunk => + translateChunkWithFallback({ + chunk, + fromLanguageCode, + toLanguageCode, + provider, + modelId, + credentials, + }), + ), ), + usageByChunk, ); translatedSegments.push(...batchResults.map(result => result.translatedVtt)); @@ -997,6 +1013,24 @@ export async function translateCaptions

, ): Promise { "use workflow"; + // Usage from provider calls made so far. A throw after translation (e.g. + // a failed S3 upload) still reports the tokens burned via the error's + // `usage` property. + const collectedUsage: TokenUsage[] = []; + try { + return await translateCaptionsInternal(assetId, trackId, toLanguageCode, options, collectedUsage); + } catch (error) { + rethrowWithTokenUsage(error, collectedUsage); + } +} + +async function translateCaptionsInternal

( + assetId: string, + trackId: string, + toLanguageCode: string, + options: TranslationOptions

, + collectedUsage: TokenUsage[], +): Promise { const { provider = "openai", model, @@ -1100,6 +1134,7 @@ export async function translateCaptions

void): unknown { + try { + fn(); + } catch (error) { + return error; + } + throw new Error("expected fn to throw"); +} + +describe("getErrorTokenUsage", () => { + it("returns undefined for non-object errors and errors without usage", () => { + expect(getErrorTokenUsage("boom")).toBeUndefined(); + expect(getErrorTokenUsage(null)).toBeUndefined(); + expect(getErrorTokenUsage(new Error("boom"))).toBeUndefined(); + }); + + it("extracts the known numeric fields from a plain usage object", () => { + const error = Object.assign(new Error("boom"), { + usage: { inputTokens: 10, outputTokens: 5, totalTokens: 15, extraneous: "x" }, + }); + + expect(getErrorTokenUsage(error)).toEqual({ + inputTokens: 10, + outputTokens: 5, + totalTokens: 15, + }); + }); + + it("falls back to nested AI SDK token details when flat fields are absent", () => { + const error = new NoObjectGeneratedError({ + finishReason: "length", + response: { + id: "resp_123", + modelId: "test-model", + timestamp: new Date("2026-03-10T00:00:00.000Z"), + }, + usage: { + inputTokens: 50, + inputTokenDetails: { + noCacheTokens: 45, + cacheReadTokens: 5, + cacheWriteTokens: undefined, + }, + outputTokens: 10, + outputTokenDetails: { + textTokens: 7, + reasoningTokens: 3, + }, + totalTokens: 60, + }, + }); + + expect(getErrorTokenUsage(error)).toEqual({ + inputTokens: 50, + outputTokens: 10, + totalTokens: 60, + reasoningTokens: 3, + cachedInputTokens: 5, + }); + }); +}); + +describe("rethrowWithTokenUsage", () => { + it("rethrows the same error with aggregate usage attached", () => { + const error = new Error("boom"); + const collected: TokenUsage[] = [ + { inputTokens: 10, outputTokens: 5, totalTokens: 15 }, + { inputTokens: 4, outputTokens: 6, totalTokens: 10 }, + ]; + + const thrown = captureThrown(() => rethrowWithTokenUsage(error, collected)); + + expect(thrown).toBe(error); + expect((thrown as { usage?: TokenUsage }).usage).toEqual({ + inputTokens: 14, + outputTokens: 11, + totalTokens: 25, + }); + }); + + it("folds the error's own usage into the aggregate", () => { + const error = Object.assign(new Error("boom"), { + usage: { inputTokens: 3, totalTokens: 3 }, + }); + + const thrown = captureThrown(() => + rethrowWithTokenUsage(error, [{ inputTokens: 10, outputTokens: 5, totalTokens: 15 }])); + + expect((thrown as { usage?: TokenUsage }).usage).toEqual({ + inputTokens: 13, + outputTokens: 5, + totalTokens: 18, + }); + }); + + it("leaves the error untouched when no usage was collected", () => { + const error = new Error("boom"); + + const thrown = captureThrown(() => rethrowWithTokenUsage(error, [])); + + expect(thrown).toBe(error); + expect(Object.hasOwn(thrown as object, "usage")).toBe(false); + }); + + it("rethrows non-object errors unchanged", () => { + const thrown = captureThrown(() => + rethrowWithTokenUsage("boom", [{ inputTokens: 1 }])); + + expect(thrown).toBe("boom"); + }); +}); + +describe("wrapError usage propagation", () => { + it("carries usage from the original error onto the wrapped error", () => { + const original = Object.assign(new Error("provider exploded"), { + usage: { inputTokens: 10, outputTokens: 5, totalTokens: 15 }, + }); + + const thrown = captureThrown(() => wrapError(original, "Failed to analyze")); + + expect(thrown).toBeInstanceOf(Error); + expect((thrown as Error).message).toBe("Failed to analyze: provider exploded"); + expect((thrown as { usage?: TokenUsage }).usage).toEqual({ + inputTokens: 10, + outputTokens: 5, + totalTokens: 15, + }); + }); + + it("does not add a usage property when the original error carries none", () => { + const thrown = captureThrown(() => wrapError(new Error("boom"), "Failed")); + + expect(Object.hasOwn(thrown as object, "usage")).toBe(false); + }); + + it("rethrows MuxAiError instances as-is, preserving any attached usage", () => { + const error = Object.assign(new MuxAiError("customer-safe"), { + usage: { totalTokens: 7 }, + }); + + const thrown = captureThrown(() => wrapError(error, "context")); + + expect(thrown).toBe(error); + expect((thrown as { usage?: TokenUsage }).usage).toEqual({ totalTokens: 7 }); + }); +}); diff --git a/tests/unit/translate-captions-error-usage.test.ts b/tests/unit/translate-captions-error-usage.test.ts new file mode 100644 index 0000000..dad9033 --- /dev/null +++ b/tests/unit/translate-captions-error-usage.test.ts @@ -0,0 +1,147 @@ +import { APICallError } from "ai"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import type { TokenUsage } from "../../src/types"; + +vi.mock("ai", async importOriginal => ({ + ...(await importOriginal()), + generateText: vi.fn(), +})); + +vi.mock("../../src/lib/mux-assets", () => ({ + getAssetDurationSecondsFromAsset: vi.fn(), + getPlaybackIdForAsset: vi.fn(), +})); + +vi.mock("../../src/lib/mux-tracks", () => ({ + createTextTrackOnMux: vi.fn(), + fetchVttFromMux: vi.fn(), +})); + +vi.mock("../../src/lib/workflow-credentials", () => ({ + resolveMuxSigningContext: vi.fn(), +})); + +vi.mock("../../src/lib/providers", async importOriginal => ({ + ...(await importOriginal()), + createLanguageModelFromConfig: vi.fn(), + resolveLanguageModelConfig: vi.fn(), +})); + +const { generateText } = await import("ai"); +const { getAssetDurationSecondsFromAsset, getPlaybackIdForAsset } = await import("../../src/lib/mux-assets"); +const { fetchVttFromMux } = await import("../../src/lib/mux-tracks"); +const { resolveMuxSigningContext } = await import("../../src/lib/workflow-credentials"); +const { createLanguageModelFromConfig, resolveLanguageModelConfig } = await import("../../src/lib/providers"); +const { translateCaptions } = await import("../../src/workflows/translate-captions"); + +// Four 30s cues so duration-based chunking (target 60s) yields two chunks of +// two cues each. +const VTT_CONTENT = [ + "WEBVTT", + "", + "1", + "00:00:00.000 --> 00:00:30.000", + "Hello one", + "", + "2", + "00:00:30.000 --> 00:01:00.000", + "Hello two", + "", + "3", + "00:01:00.000 --> 00:01:30.000", + "Hello three", + "", + "4", + "00:01:30.000 --> 00:02:00.000", + "Hello four", + "", +].join("\n"); + +const CHUNK_USAGE = { inputTokens: 100, outputTokens: 50, totalTokens: 150 }; + +const TRANSLATION_OPTIONS = { + provider: "openai" as const, + uploadToS3: false, + uploadToMux: false, + chunking: { + enabled: true, + minimumAssetDurationSeconds: 1, + targetChunkDurationSeconds: 60, + maxConcurrentTranslations: 1, + }, +}; + +async function captureRejection(promise: Promise): Promise { + try { + await promise; + } catch (error) { + return error; + } + throw new Error("expected promise to reject"); +} + +beforeEach(() => { + vi.resetAllMocks(); + vi.spyOn(console, "warn").mockImplementation(() => {}); + + vi.mocked(getPlaybackIdForAsset).mockResolvedValue({ + asset: { + id: "asset-123", + tracks: [ + { id: "track-1", type: "text", status: "ready", language_code: "en", text_type: "subtitles" }, + ], + }, + playbackId: "playback-123", + policy: "public", + } as any); + vi.mocked(getAssetDurationSecondsFromAsset).mockReturnValue(120); + vi.mocked(resolveMuxSigningContext).mockResolvedValue(undefined); + vi.mocked(fetchVttFromMux).mockResolvedValue(VTT_CONTENT); + vi.mocked(resolveLanguageModelConfig).mockReturnValue({ provider: "openai", modelId: "gpt-test" } as any); + vi.mocked(createLanguageModelFromConfig).mockResolvedValue({} as any); +}); + +describe("translateCaptions error-path token usage", () => { + it("attaches the usage of completed chunks when a later chunk fails", async () => { + vi.mocked(generateText).mockImplementation((async (args: any) => { + const userContent = args.messages.find((m: any) => m.role === "user").content as string; + if (userContent.includes("Hello three")) { + // Non-retryable provider error → the chunk fails fast without + // splitting, after the first chunk already burned tokens. + throw new APICallError({ + message: "Bad request", + requestBodyValues: {}, + statusCode: 400, + url: "https://api.example.test/v1/messages", + }); + } + const cueCount = Number(/Return exactly (\d+) translated/.exec(userContent)?.[1] ?? 0); + const translations = Array.from({ length: cueCount }, (_, i) => `hola ${i}`); + return { + output: { translations }, + text: JSON.stringify({ translations }), + usage: CHUNK_USAGE, + }; + }) as any); + + const error = await captureRejection( + translateCaptions("asset-123", "track-1", "es", TRANSLATION_OPTIONS), + ); + + expect((error as Error).message).toContain("Failed to translate VTT"); + // One chunk of two completed before the failure, so exactly one chunk's + // usage must ride on the error. + expect((error as { usage?: TokenUsage }).usage).toEqual(CHUNK_USAGE); + }); + + it("attaches no usage when the workflow fails before any provider call", async () => { + const error = await captureRejection( + translateCaptions("asset-123", "missing-track", "es", TRANSLATION_OPTIONS), + ); + + expect((error as Error).message).toContain("missing-track"); + expect(Object.hasOwn(error as object, "usage")).toBe(false); + expect(vi.mocked(generateText)).not.toHaveBeenCalled(); + }); +}); From e477849cb7898919049dd90681c2e3c20c5357b2 Mon Sep 17 00:00:00 2001 From: Walker Date: Tue, 14 Jul 2026 10:09:29 -0700 Subject: [PATCH 2/2] cursor --- src/lib/mux-ai-error.ts | 5 +---- src/lib/token-usage.ts | 25 ++++++++----------------- src/workflows/ask-questions.ts | 3 --- src/workflows/burned-in-captions.ts | 2 -- src/workflows/chapters.ts | 3 --- src/workflows/edit-captions.ts | 3 --- src/workflows/embeddings.ts | 14 +++++++------- src/workflows/engagement-insights.ts | 3 --- src/workflows/summarization.ts | 3 --- src/workflows/translate-captions.ts | 3 --- 10 files changed, 16 insertions(+), 48 deletions(-) diff --git a/src/lib/mux-ai-error.ts b/src/lib/mux-ai-error.ts index 5d4cb68..e7ef26f 100644 --- a/src/lib/mux-ai-error.ts +++ b/src/lib/mux-ai-error.ts @@ -83,10 +83,7 @@ export function wrapError(error: unknown, message: string): never { console.warn(`[@mux/ai] Suppressed suspected prompt leak in wrapped error (context: ${message}, reason: ${leakReason}).`); } const wrapped = new Error(`${message}: ${detail}`); - // Token usage rides on errors as a plain `usage` property (AI SDK errors - // like NoObjectGeneratedError, plus errors annotated by - // rethrowWithTokenUsage). Carry it through wrapping so failed workflows - // can still report the tokens they burned. + // Carry `usage` through wrapping so failed workflows still report tokens burned. const usage = getErrorTokenUsage(error); if (usage) { (wrapped as Error & { usage?: TokenUsage }).usage = usage; diff --git a/src/lib/token-usage.ts b/src/lib/token-usage.ts index 3129b4c..48220de 100644 --- a/src/lib/token-usage.ts +++ b/src/lib/token-usage.ts @@ -33,14 +33,10 @@ export function aggregateTokenUsage(usages: TokenUsage[]): TokenUsage { } /** - * Reads token usage carried on a thrown error's plain `usage` property. - * - * Covers both AI SDK errors that report usage for the failed call (e.g. - * `NoObjectGeneratedError`) and errors this package has already annotated - * via {@link rethrowWithTokenUsage}. Only the known numeric fields are - * extracted; AI SDK v6 nests reasoning/cache counts under - * `outputTokenDetails`/`inputTokenDetails` when the deprecated flat fields - * are absent, so those are used as fallbacks. + * 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) { @@ -77,15 +73,10 @@ export function getErrorTokenUsage(error: unknown): TokenUsage | undefined { } /** - * Rethrows `error` with the aggregate token usage of all provider calls made - * so far attached as a plain enumerable `usage` property, so consumers can - * report tokens burned by failed workflows (mirrors the AI SDK's - * `NoObjectGeneratedError.usage` convention). - * - * The error's own `usage` (from the failing call) is folded into the - * aggregate. When no usage was collected and the error carries none, the - * error is rethrown untouched — callers treat a missing `usage` as "no - * tokens spent". + * 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]; diff --git a/src/workflows/ask-questions.ts b/src/workflows/ask-questions.ts index 2da508c..2772716 100644 --- a/src/workflows/ask-questions.ts +++ b/src/workflows/ask-questions.ts @@ -755,9 +755,6 @@ export async function askQuestions( options?: AskQuestionsOptions, ): Promise { "use workflow"; - // Usage from provider calls made so far. A throw after the analysis call - // (e.g. incomplete answers) still reports the tokens burned via the - // error's `usage` property. const collectedUsage: TokenUsage[] = []; try { return await askQuestionsInternal(assetId, questions, options, collectedUsage); diff --git a/src/workflows/burned-in-captions.ts b/src/workflows/burned-in-captions.ts index c1f12fd..0f9fb45 100644 --- a/src/workflows/burned-in-captions.ts +++ b/src/workflows/burned-in-captions.ts @@ -329,8 +329,6 @@ export async function hasBurnedInCaptions( options: BurnedInCaptionsOptions = {}, ): Promise { "use workflow"; - // Usage from provider calls made so far. A throw after the analysis call - // still reports the tokens burned via the error's `usage` property. const collectedUsage: TokenUsage[] = []; try { return await hasBurnedInCaptionsInternal(assetId, options, collectedUsage); diff --git a/src/workflows/chapters.ts b/src/workflows/chapters.ts index f5adc12..0c76d1c 100644 --- a/src/workflows/chapters.ts +++ b/src/workflows/chapters.ts @@ -369,9 +369,6 @@ export async function generateChapters( options: ChaptersOptions = {}, ): Promise { "use workflow"; - // Usage from provider calls made so far. A throw after the generation call - // (e.g. no valid chapters) still reports the tokens burned via the error's - // `usage` property. const collectedUsage: TokenUsage[] = []; try { return await generateChaptersInternal(assetId, options, collectedUsage); diff --git a/src/workflows/edit-captions.ts b/src/workflows/edit-captions.ts index 5349353..b47c1ca 100644 --- a/src/workflows/edit-captions.ts +++ b/src/workflows/edit-captions.ts @@ -551,9 +551,6 @@ export async function editCaptions

, ): Promise { "use workflow"; - // Usage from provider calls made so far. A throw after profanity detection - // (e.g. a failed S3 upload) still reports the tokens burned via the - // error's `usage` property. const collectedUsage: TokenUsage[] = []; try { return await editCaptionsInternal(assetId, trackId, options, collectedUsage); diff --git a/src/workflows/embeddings.ts b/src/workflows/embeddings.ts index adc7699..71cb75b 100644 --- a/src/workflows/embeddings.ts +++ b/src/workflows/embeddings.ts @@ -7,7 +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 { rethrowWithTokenUsage } from "../lib/token-usage.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"; @@ -114,9 +114,6 @@ async function generateSingleChunkEmbedding({ tokenCount: chunk.tokenCount, }, }, - // Embedding tokens are all input tokens. Reported so a failure partway - // through the batches can attach the tokens burned so far to the error; - // the success-path result intentionally remains unchanged. usage: typeof response.usage?.tokens === "number" ? { inputTokens: response.usage.tokens, totalTokens: response.usage.tokens } : undefined, @@ -223,9 +220,8 @@ async function generateEmbeddingsInternal( ), ); - // Record usage from every settled call — including the fulfilled - // siblings of a failed call — before rethrowing, so the tokens - // burned so far ride on the thrown error. + // 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") { @@ -234,6 +230,10 @@ async function generateEmbeddingsInternal( } chunkEmbeddings.push(outcome.value.chunk); } else { + const failureUsage = getErrorTokenUsage(outcome.reason); + if (failureUsage) { + collectedUsage.push(failureUsage); + } failures.push(outcome.reason); } } diff --git a/src/workflows/engagement-insights.ts b/src/workflows/engagement-insights.ts index 3e37cf4..5e9f04d 100644 --- a/src/workflows/engagement-insights.ts +++ b/src/workflows/engagement-insights.ts @@ -643,9 +643,6 @@ export async function generateEngagementInsights( options: EngagementInsightsOptions = {}, ): Promise { "use workflow"; - // Usage from provider calls made so far. A throw after the insights call - // (e.g. no valid insights) still reports the tokens burned via the error's - // `usage` property. const collectedUsage: TokenUsage[] = []; try { return await generateEngagementInsightsInternal(assetId, options, collectedUsage); diff --git a/src/workflows/summarization.ts b/src/workflows/summarization.ts index 0b6ebf2..35ef5c4 100644 --- a/src/workflows/summarization.ts +++ b/src/workflows/summarization.ts @@ -684,9 +684,6 @@ export async function getSummaryAndTags( options?: SummarizationOptions, ): Promise { "use workflow"; - // Usage from provider calls made so far. A throw after the analysis call - // (e.g. unusable model output) still reports the tokens burned via the - // error's `usage` property. const collectedUsage: TokenUsage[] = []; try { return await getSummaryAndTagsInternal(assetId, options, collectedUsage); diff --git a/src/workflows/translate-captions.ts b/src/workflows/translate-captions.ts index 3120841..b71aa3c 100644 --- a/src/workflows/translate-captions.ts +++ b/src/workflows/translate-captions.ts @@ -1013,9 +1013,6 @@ export async function translateCaptions

, ): Promise { "use workflow"; - // Usage from provider calls made so far. A throw after translation (e.g. - // a failed S3 upload) still reports the tokens burned via the error's - // `usage` property. const collectedUsage: TokenUsage[] = []; try { return await translateCaptionsInternal(assetId, trackId, toLanguageCode, options, collectedUsage);