Feat/retry count

packages/scrawn/src/core/ai/wrap.ts

@@ -68,26 +68,30 @@ export function createBillableAI<TTag extends string>(
       }) => {
         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<string, number | undefined>)
+                ?.inputTokens,
+              completionTokens: (
+                event.usage as Record<string, number | undefined>
+              )?.outputTokens,
+              totalTokens: (event.usage as Record<string, number | undefined>)
+                ?.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

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<TTags>,
-    defaults: {
-      inputDebit: Debit<TTags>;
-      outputDebit: Debit<TTags>;
-      inputCacheDebit: Debit<TTags>;
-      outputCacheDebit: Debit<TTags>;
-      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<TTags>;
+    outputDebit: Debit<TTags>;
+    inputCacheDebit?: Debit<TTags>;
+    outputCacheDebit?: Debit<TTags>;
+    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;

packages/scrawn/src/core/types/event.ts

@@ -243,6 +243,8 @@ export type PayloadExtractor<TTag extends string = string> = (
  * 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<void>;
 }