diff --git a/packages/scrawn/src/core/ai/wrap.ts b/packages/scrawn/src/core/ai/wrap.ts index 509d08d..7848a08 100644 --- a/packages/scrawn/src/core/ai/wrap.ts +++ b/packages/scrawn/src/core/ai/wrap.ts @@ -68,26 +68,30 @@ export function createBillableAI( }) => { if (!event.usage) return; - biller.trackAI( + biller.trackAI({ userId, - { - modelId: event.model?.modelId ?? "unknown", - provider: event.model?.provider ?? "unknown", + event: { + model: { + modelId: event.model?.modelId ?? "unknown", + provider: event.model?.provider ?? "unknown", + }, + usage: { + promptTokens: (event.usage as Record) + ?.inputTokens, + completionTokens: ( + event.usage as Record + )?.outputTokens, + totalTokens: (event.usage as Record) + ?.totalTokens, + }, }, - { - inputTokens: (event.usage.inputTokens as number) ?? 0, - outputTokens: (event.usage.outputTokens as number) ?? 0, - totalTokens: (event.usage.totalTokens as number) ?? 0, - inputCachedTokens: event.usage.inputCachedTokens as - | number - | undefined, - outputCachedTokens: event.usage.outputCachedTokens as - | number - | undefined, - }, - billing, - defaults - ); + inputDebit: billing.inputDebit ?? defaults.inputDebit, + outputDebit: billing.outputDebit ?? defaults.outputDebit, + inputCacheDebit: billing.inputCacheDebit ?? defaults.inputCacheDebit, + outputCacheDebit: + billing.outputCacheDebit ?? defaults.outputCacheDebit, + provider: billing.provider ?? defaults.provider, + }); }; // Chain billing + user callbacks diff --git a/packages/scrawn/src/core/scrawn.ts b/packages/scrawn/src/core/scrawn.ts index 074ce1c..39ff19e 100644 --- a/packages/scrawn/src/core/scrawn.ts +++ b/packages/scrawn/src/core/scrawn.ts @@ -43,6 +43,7 @@ import { type CreateCheckoutLinkResponse, } from "../gen/payment/v1/payment.js"; import { + ScrawnError, ScrawnConfigError, ScrawnValidationError, convertGrpcError, @@ -492,8 +493,27 @@ export class Scrawn< ? (error as import("./errors/index.js").ScrawnError) : convertGrpcError(error); + let manualRetryCount = 0; + const maxManualRetries = this.retryCount; + const retryContext: RetryContext = { + get retryCount() { + return manualRetryCount; + }, retry: async () => { + if (manualRetryCount >= maxManualRetries) { + const exceededError = new ScrawnError( + "Manual retry limit exceeded", + { + code: "RETRY_LIMIT_EXCEEDED", + retryable: false, + details: { retriesAttempted: manualRetryCount }, + } + ); + options.onError!(exceededError); + return; + } + manualRetryCount++; try { await attempt(); } catch (retryError) { @@ -1177,45 +1197,81 @@ export class Scrawn< /** * Manually track AI token usage from an event callback. * - * Converts a Vercel AI SDK `onStepFinish` or `onFinish` event into - * an `AITokenUsagePayload` and streams it to the backend (fire-and-forget). + * Accepts the full event object from `onStepFinish` or `onFinish` + * and extracts `model` + `usage` automatically. * * Use this for manual control when you don't want the full `biller.ai()` wrapper. * - * @param userId - The user ID to bill against - * @param model - Model info (modelId + provider from the event) - * @param usage - Token usage from the event (event.usage or event.totalUsage) - * @param overrides - Override billing per-call (optional) - * @param defaults - Fallback debit config (required — use same as biller.ai opts) - * * @example * ```typescript * const result = await ai.streamText({ * model: openai("gpt-4o"), * prompt: "Hello", * onStepFinish: event => { - * biller.trackAI("user-123", event.model, event.usage, {}, { - * inputDebit: { tag: "AI_INPUT" }, - * outputDebit: { tag: "AI_OUTPUT" }, + * biller.trackAI({ + * userId: "user-123", + * event, + * inputDebit: biller.tag("AI_INPUT"), + * outputDebit: biller.tag("AI_OUTPUT"), * }); * }, * }); * ``` */ - trackAI( - userId: string, - model: ModelInfo, - usage: LanguageModelUsage, - overrides: BillableCallParams, - defaults: { - inputDebit: Debit; - outputDebit: Debit; - inputCacheDebit: Debit; - outputCacheDebit: Debit; - provider?: string; - } - ): void { - const payload = buildAIPayload(userId, model, usage, overrides, defaults); + trackAI(config: { + userId: string; + event: { + model: ModelInfo; + usage?: { + promptTokens?: number; + completionTokens?: number; + totalTokens?: number; + }; + totalUsage?: { + promptTokens?: number; + completionTokens?: number; + totalTokens?: number; + }; + }; + inputDebit: Debit; + outputDebit: Debit; + inputCacheDebit?: Debit; + outputCacheDebit?: Debit; + provider?: string; + }): void { + const { + userId, + event, + inputDebit, + outputDebit, + inputCacheDebit, + outputCacheDebit, + provider, + } = config; + const usage = event.usage ?? event.totalUsage ?? {}; + const model: ModelInfo = event.model; + + const mappedUsage: LanguageModelUsage = { + inputTokens: usage.promptTokens ?? 0, + outputTokens: usage.completionTokens ?? 0, + totalTokens: usage.totalTokens ?? 0, + inputCachedTokens: 0, + outputCachedTokens: 0, + }; + + const payload = buildAIPayload( + userId, + model, + mappedUsage, + { inputDebit, outputDebit, inputCacheDebit, outputCacheDebit, provider }, + { + inputDebit, + outputDebit, + inputCacheDebit: inputCacheDebit ?? inputDebit, + outputCacheDebit: outputCacheDebit ?? outputDebit, + provider, + } + ); this.aiTokenStreamConsumer( (async function* () { yield payload; diff --git a/packages/scrawn/src/core/types/event.ts b/packages/scrawn/src/core/types/event.ts index 7c8a9ce..64fca62 100644 --- a/packages/scrawn/src/core/types/event.ts +++ b/packages/scrawn/src/core/types/event.ts @@ -243,6 +243,8 @@ export type PayloadExtractor = ( * Only provided when the method supports manual retry (basicUsageEventConsumer). */ export interface RetryContext { + /** Number of manual retries attempted so far (starts at 0). */ + retryCount: number; /** Re-attempt the failed operation using the same eventId and idempotencyKey. */ retry: () => Promise; }