From dc6e632ce518a018f31e5e3ca7ccb08372e9c741 Mon Sep 17 00:00:00 2001
From: Devyash Saini <dysaini2004@gmail.com>
Date: Tue, 26 May 2026 22:02:40 +0530
Subject: [PATCH 01/11] feat(core): add biller.ai() Vercel AI SDK wrapper with
 auto-billing

- biller.ai(sdk, opts) returns proxied AI SDK with per-step billing
- Auto-injects onStepFinish for streamText, generateText, streamObject, generateObject
- Chains user's own callbacks alongside billing (non-blocking)
- biller.trackAI() for manual tracking at onStepFinish/onFinish
- Supports cached input/output token billing with separate debit config
- Per-call overrides for inputDebit/outputDebit/cacheDebit
- Zero any types, fully generic over TTags
---
 packages/scrawn/src/core/ai/track.ts |  45 ++++++++
 packages/scrawn/src/core/ai/types.ts |  71 +++++++++++++
 packages/scrawn/src/core/ai/wrap.ts  | 151 +++++++++++++++++++++++++++
 packages/scrawn/src/core/scrawn.ts   |  96 +++++++++++++++++
 packages/scrawn/src/index.ts         |   8 ++
 5 files changed, 371 insertions(+)
 create mode 100644 packages/scrawn/src/core/ai/track.ts
 create mode 100644 packages/scrawn/src/core/ai/types.ts
 create mode 100644 packages/scrawn/src/core/ai/wrap.ts

diff --git a/packages/scrawn/src/core/ai/track.ts b/packages/scrawn/src/core/ai/track.ts
new file mode 100644
index 0000000..ec3a4b2
--- /dev/null
+++ b/packages/scrawn/src/core/ai/track.ts
@@ -0,0 +1,45 @@
+import type { AITokenUsagePayload, DebitField } from "../types/event.js";
+import type {
+  BillableCallParams,
+  LanguageModelUsage,
+  ModelInfo,
+} from "./types.js";
+
+/**
+ * Builds an AITokenUsagePayload from an AI SDK step/finish event.
+ * Falls back to regular debit pricing for cache tokens if not specified.
+ */
+export function buildAIPayload<TTag extends string = string>(
+  userId: string,
+  model: ModelInfo,
+  usage: LanguageModelUsage,
+  overrides: BillableCallParams<TTag>,
+  defaults: {
+    inputDebit: DebitField<TTag>;
+    outputDebit: DebitField<TTag>;
+    inputCacheDebit: DebitField<TTag>;
+    outputCacheDebit: DebitField<TTag>;
+    provider?: string;
+  }
+): AITokenUsagePayload<TTag> {
+  return {
+    userId,
+    model: model.modelId,
+    inputTokens: usage.inputTokens,
+    outputTokens: usage.outputTokens,
+    inputDebit: overrides.inputDebit ?? defaults.inputDebit,
+    outputDebit: overrides.outputDebit ?? defaults.outputDebit,
+    provider:
+      (overrides.provider as string) ?? defaults.provider ?? model.provider,
+    inputCacheTokens:
+      overrides.inputCacheDebit ?? usage.inputCachedTokens !== undefined
+        ? usage.inputCachedTokens
+        : undefined,
+    inputCacheDebit: overrides.inputCacheDebit ?? defaults.inputCacheDebit,
+    outputCacheTokens:
+      overrides.outputCacheDebit ?? usage.outputCachedTokens !== undefined
+        ? usage.outputCachedTokens
+        : undefined,
+    outputCacheDebit: overrides.outputCacheDebit ?? defaults.outputCacheDebit,
+  };
+}
diff --git a/packages/scrawn/src/core/ai/types.ts b/packages/scrawn/src/core/ai/types.ts
new file mode 100644
index 0000000..6323abc
--- /dev/null
+++ b/packages/scrawn/src/core/ai/types.ts
@@ -0,0 +1,71 @@
+import type { DebitField } from "../types/event.js";
+
+/**
+ * Configuration for the biller.ai() wrapper.
+ * Default billing settings applied automatically to all AI SDK calls.
+ */
+export interface BillableAIOptions<TTag extends string = string> {
+  /** Default billing for input tokens (required). */
+  inputDebit: DebitField<TTag>;
+  /** Default billing for output tokens (required). */
+  outputDebit: DebitField<TTag>;
+  /** Default billing for cached input tokens. Falls back to inputDebit if not set. */
+  inputCacheDebit?: DebitField<TTag>;
+  /** Default billing for cached output tokens. Falls back to outputDebit if not set. */
+  outputCacheDebit?: DebitField<TTag>;
+  /** Default provider override. If not set, auto-detected from the model's provider. */
+  provider?: string;
+}
+
+/**
+ * Additional fields injected by the AI SDK wrapper into function params.
+ * All original AI SDK params are preserved; only userId is added.
+ */
+export interface BillableCallParams<TTag extends string = string> {
+  /** The user ID to bill against. */
+  userId: string;
+  /** Override input token billing for this specific call. */
+  inputDebit?: DebitField<TTag>;
+  /** Override output token billing for this specific call. */
+  outputDebit?: DebitField<TTag>;
+  /** Override cached input token billing for this specific call. */
+  inputCacheDebit?: DebitField<TTag>;
+  /** Override cached output token billing for this specific call. */
+  outputCacheDebit?: DebitField<TTag>;
+  /** Override provider for this specific call. */
+  provider?: string;
+}
+
+/**
+ * Strips the BillableCallParams from a type, leaving only original AI SDK params.
+ * Used internally to forward the original params to the real AI SDK.
+ */
+export type StripScrawnParams<T> = Omit<T, keyof BillableCallParams>;
+
+/**
+ * Language model usage as returned by Vercel AI SDK event listeners.
+ * Matches the shape of OnStepFinishEvent.usage and OnFinishEvent.totalUsage.
+ */
+export interface LanguageModelUsage {
+  /** Number of input (prompt) tokens consumed. */
+  inputTokens: number;
+  /** Number of output (completion) tokens consumed. */
+  outputTokens: number;
+  /** Total tokens consumed (inputTokens + outputTokens). */
+  totalTokens: number;
+  /** Cached input tokens (e.g., prompt caching). */
+  inputCachedTokens?: number;
+  /** Cached output tokens. */
+  outputCachedTokens?: number;
+}
+
+/**
+ * Minimal subset of the AI SDK model info needed for billing.
+ * Comes from OnStepFinishEvent.model or OnFinishEvent.model.
+ */
+export interface ModelInfo {
+  /** Model ID, e.g. "gpt-4o-mini". */
+  modelId: string;
+  /** Provider name, e.g. "openai", "anthropic". */
+  provider: string;
+}
diff --git a/packages/scrawn/src/core/ai/wrap.ts b/packages/scrawn/src/core/ai/wrap.ts
new file mode 100644
index 0000000..8f0169a
--- /dev/null
+++ b/packages/scrawn/src/core/ai/wrap.ts
@@ -0,0 +1,151 @@
+import type { Scrawn } from "../scrawn.js";
+import type {
+  BillableAIOptions,
+  BillableCallParams,
+  StripScrawnParams,
+  ModelInfo,
+} from "./types.js";
+import { buildAIPayload } from "./track.js";
+
+/** AI SDK function names that accept event callbacks and should be wrapped. */
+const BILLABLE_FNS = [
+  "streamText",
+  "generateText",
+  "streamObject",
+  "generateObject",
+] as const;
+
+type BillableFnName = (typeof BILLABLE_FNS)[number];
+
+/** An AI SDK module shape — duck-typed for flexibility. */
+type AISDKModule = {
+  [K in BillableFnName]?: (...args: any[]) => Promise<unknown>;
+};
+
+/**
+ * Returns a proxied AI SDK module. Each text generation function is wrapped to:
+ * 1. Accept a `userId` field (and optional billing overrides)
+ * 2. Auto-inject `onStepFinish` to track billing on every step
+ * 3. Chain the user's own `onStepFinish`/`onFinish` after billing
+ *
+ * The returned object has the same types as the original AI SDK,
+ * with the billable params injected.
+ */
+export function createBillableAI<TTag extends string>(
+  sdk: AISDKModule,
+  biller: Scrawn<TTag>,
+  opts: BillableAIOptions<TTag>
+): Record<string, unknown> {
+  const proxied: Record<string, unknown> = { ...sdk };
+
+  for (const fnName of BILLABLE_FNS) {
+    const original = sdk[fnName as BillableFnName];
+    if (typeof original !== "function") continue;
+
+    proxied[fnName] = (...args: unknown[]): unknown => {
+      const params = (args[0] ?? {}) as Record<string, unknown>;
+      const userId = params.userId as string | undefined;
+      const billing: BillableCallParams<TTag> = extractBillingParams(params);
+
+      if (userId === undefined || userId.trim() === "") {
+        // No userId — pass through to original unchanged
+        return original.apply(sdk, args);
+      }
+
+      const { onStepFinish: userStep, onFinish: userFinish, ...rest } = params;
+      const billingParams = { ...rest };
+
+      // Inject onStepFinish for per-step billing
+      const billingStep = (event: {
+        model: ModelInfo;
+        usage: Record<string, unknown>;
+      }) => {
+        if (!event.usage) return;
+
+        const payload = buildAIPayload(
+          userId,
+          {
+            modelId: event.model?.modelId ?? "unknown",
+            provider: event.model?.provider ?? "unknown",
+          },
+          {
+            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,
+          {
+            inputDebit: opts.inputDebit,
+            outputDebit: opts.outputDebit,
+            inputCacheDebit: opts.inputCacheDebit ?? opts.inputDebit,
+            outputCacheDebit: opts.outputCacheDebit ?? opts.outputDebit,
+            provider: opts.provider,
+          }
+        );
+
+        biller.aiTokenStreamConsumer(
+          (async function* () {
+            yield payload;
+          })()
+        );
+      };
+
+      // Chain billing + user callbacks
+      if (
+        typeof userStep === "function" ||
+        userStep === undefined ||
+        userStep === null
+      ) {
+        billingParams.onStepFinish = chainHandlers(
+          billingStep,
+          userStep as ((e: unknown) => void) | undefined
+        );
+      }
+      if (typeof userFinish === "function") {
+        billingParams.onFinish = (event: unknown) => {
+          billingStep(
+            event as { model: ModelInfo; usage: Record<string, unknown> }
+          );
+          (userFinish as (e: unknown) => void)(event);
+        };
+      }
+
+      return original.call(sdk, billingParams);
+    };
+  }
+
+  return proxied;
+}
+
+function extractBillingParams<TTag extends string>(
+  params: Record<string, unknown>
+): BillableCallParams<TTag> {
+  return {
+    userId: "",
+    inputDebit: params.inputDebit as BillableCallParams<TTag>["inputDebit"],
+    outputDebit: params.outputDebit as BillableCallParams<TTag>["outputDebit"],
+    inputCacheDebit:
+      params.inputCacheDebit as BillableCallParams<TTag>["inputCacheDebit"],
+    outputCacheDebit:
+      params.outputCacheDebit as BillableCallParams<TTag>["outputCacheDebit"],
+  };
+}
+
+function chainHandlers(
+  first: (e: { model: ModelInfo; usage: Record<string, unknown> }) => void,
+  second:
+    | ((e: { model: ModelInfo; usage: Record<string, unknown> }) => void)
+    | undefined
+): (e: { model: ModelInfo; usage: Record<string, unknown> }) => void {
+  if (!second) return first;
+  return (e: { model: ModelInfo; usage: Record<string, unknown> }) => {
+    first(e);
+    second(e);
+  };
+}
diff --git a/packages/scrawn/src/core/scrawn.ts b/packages/scrawn/src/core/scrawn.ts
index 93c8f36..cb5c425 100644
--- a/packages/scrawn/src/core/scrawn.ts
+++ b/packages/scrawn/src/core/scrawn.ts
@@ -8,6 +8,7 @@ import type {
   AITokenUsagePayload,
   EventConsumerErrorCallback,
   RetryContext,
+  DebitField,
 } from "./types/event.js";
 import type {
   AuthRegistry,
@@ -59,6 +60,14 @@ import {
   prettyPrintExpr,
   tag as _tag,
 } from "./pricing/index.js";
+import { createBillableAI } from "./ai/wrap.js";
+import type { BillableAIOptions } from "./ai/types.js";
+import { buildAIPayload } from "./ai/track.js";
+import type {
+  LanguageModelUsage,
+  ModelInfo,
+  BillableCallParams,
+} from "./ai/types.js";
 import { ScrawnConfig } from "../config.js";
 import { randomUUID } from "node:crypto";
 
@@ -1180,6 +1189,93 @@ export class Scrawn<
       yield request;
     }
   }
+
+  /**
+   * Wraps the Vercel AI SDK with automatic per-step billing.
+   *
+   * Returns a proxied version of the AI SDK where `streamText`, `generateText`,
+   * `streamObject`, and `generateObject` automatically:
+   * 1. Accept a `userId` parameter for billing
+   * 2. Fire-and-forget `aiTokenStreamConsumer` on every step via `onStepFinish`
+   * 3. Chain the user's own `onStepFinish`/`onFinish` callbacks
+   *
+   * @param sdk - The Vercel AI SDK module (import * as ai from "ai")
+   * @param opts - Default billing configuration for all calls
+   * @returns A proxied AI SDK with billing injected
+   *
+   * @example
+   * ```typescript
+   * import * as ai from "ai";
+   * import { biller } from "./scrawn/biller";
+   *
+   * const aii = biller.ai(ai, {
+   *   inputDebit: { tag: "AI_INPUT" },
+   *   outputDebit: { tag: "AI_OUTPUT" },
+   * });
+   *
+   * const result = await aii.streamText({
+   *   userId: "user-123",
+   *   model: openai("gpt-4o-mini"),
+   *   prompt: "Write a story.",
+   *   onStepFinish: event => { console.log("Step done"); },
+   * });
+   * ```
+   */
+  ai(
+    sdk: Record<string, unknown>,
+    opts: BillableAIOptions<TTags>
+  ): Record<string, unknown> {
+    return createBillableAI(sdk, this, opts);
+  }
+
+  /**
+   * 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).
+   *
+   * 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" },
+   *     });
+   *   },
+   * });
+   * ```
+   */
+  trackAI(
+    userId: string,
+    model: ModelInfo,
+    usage: LanguageModelUsage,
+    overrides: BillableCallParams<TTags>,
+    defaults: {
+      inputDebit: DebitField<TTags>;
+      outputDebit: DebitField<TTags>;
+      inputCacheDebit: DebitField<TTags>;
+      outputCacheDebit: DebitField<TTags>;
+      provider?: string;
+    }
+  ): void {
+    const payload = buildAIPayload(userId, model, usage, overrides, defaults);
+    this.aiTokenStreamConsumer(
+      (async function* () {
+        yield payload;
+      })()
+    );
+  }
 }
 
 /**
diff --git a/packages/scrawn/src/index.ts b/packages/scrawn/src/index.ts
index ef994f3..4eb769d 100644
--- a/packages/scrawn/src/index.ts
+++ b/packages/scrawn/src/index.ts
@@ -61,3 +61,11 @@ export type {
 // Export central configuration
 export { ScrawnConfig, scrawnConfig } from "./config.js";
 export type { ScrawnCLIConfig } from "./config.js";
+
+// Export AI SDK wrapper types
+export type {
+  BillableAIOptions,
+  BillableCallParams,
+  LanguageModelUsage,
+  ModelInfo,
+} from "./core/ai/types.ts";

From 58528319ae5594ce5fcb2d79c3a80e8124eac016 Mon Sep 17 00:00:00 2001
From: Devyash Saini <dysaini2004@gmail.com>
Date: Tue, 26 May 2026 22:09:23 +0530
Subject: [PATCH 02/11] refactor(ai): use biller.trackAI() in wrap instead of
 duplicate payload logic

- wrap.ts billingStep now delegates to biller.trackAI()
- Removes duplicate buildAIPayload + aiTokenStreamConsumer calls
- trackAI is the single canonical billing entry point
---
 packages/scrawn/src/core/ai/wrap.ts | 27 +++++++++++----------------
 1 file changed, 11 insertions(+), 16 deletions(-)

diff --git a/packages/scrawn/src/core/ai/wrap.ts b/packages/scrawn/src/core/ai/wrap.ts
index 8f0169a..a1e6bbb 100644
--- a/packages/scrawn/src/core/ai/wrap.ts
+++ b/packages/scrawn/src/core/ai/wrap.ts
@@ -5,7 +5,6 @@ import type {
   StripScrawnParams,
   ModelInfo,
 } from "./types.js";
-import { buildAIPayload } from "./track.js";
 
 /** AI SDK function names that accept event callbacks and should be wrapped. */
 const BILLABLE_FNS = [
@@ -55,14 +54,22 @@ export function createBillableAI<TTag extends string>(
       const { onStepFinish: userStep, onFinish: userFinish, ...rest } = params;
       const billingParams = { ...rest };
 
-      // Inject onStepFinish for per-step billing
+      const defaults = {
+        inputDebit: opts.inputDebit,
+        outputDebit: opts.outputDebit,
+        inputCacheDebit: opts.inputCacheDebit ?? opts.inputDebit,
+        outputCacheDebit: opts.outputCacheDebit ?? opts.outputDebit,
+        provider: opts.provider,
+      };
+
+      // Inject onStepFinish for per-step billing using biller.trackAI
       const billingStep = (event: {
         model: ModelInfo;
         usage: Record<string, unknown>;
       }) => {
         if (!event.usage) return;
 
-        const payload = buildAIPayload(
+        biller.trackAI(
           userId,
           {
             modelId: event.model?.modelId ?? "unknown",
@@ -80,19 +87,7 @@ export function createBillableAI<TTag extends string>(
               | undefined,
           },
           billing,
-          {
-            inputDebit: opts.inputDebit,
-            outputDebit: opts.outputDebit,
-            inputCacheDebit: opts.inputCacheDebit ?? opts.inputDebit,
-            outputCacheDebit: opts.outputCacheDebit ?? opts.outputDebit,
-            provider: opts.provider,
-          }
-        );
-
-        biller.aiTokenStreamConsumer(
-          (async function* () {
-            yield payload;
-          })()
+          defaults
         );
       };
 

From 2762c9f6e971a4063f937397c804ac90a4211698 Mon Sep 17 00:00:00 2001
From: Devyash Saini <dysaini2004@gmail.com>
Date: Tue, 26 May 2026 22:15:03 +0530
Subject: [PATCH 03/11] feat(examples): add ai-sdk-wrapper-usage.ts
 demonstrating biller.ai()

- Level 1: biller.ai(ai, {inputDebit, outputDebit}) auto-wrapper
- Level 1b: user's own onStepFinish chained with billing
- Level 2: Manual biller.trackAI() in onStepFinish callback
- Demonstrates both streamText and generateText patterns
---
 examples/ai-sdk-wrapper-usage.ts | 92 ++++++++++++++++++++++++++++++++
 examples/bun.lock                | 27 ++++++++++
 examples/package.json            |  8 +--
 3 files changed, 124 insertions(+), 3 deletions(-)
 create mode 100644 examples/ai-sdk-wrapper-usage.ts

diff --git a/examples/ai-sdk-wrapper-usage.ts b/examples/ai-sdk-wrapper-usage.ts
new file mode 100644
index 0000000..223a3d2
--- /dev/null
+++ b/examples/ai-sdk-wrapper-usage.ts
@@ -0,0 +1,92 @@
+import * as ai from "ai";
+import { openai } from "@ai-sdk/openai";
+import { biller } from "./scrawn/biller.js";
+import { config as dotenvConfig } from "dotenv";
+dotenvConfig({ path: ".env.local" });
+
+// Helper: cast the wrapped function back to a usable callable type
+type WrappedFn = (
+  params: Record<string, unknown>
+) => Promise<Record<string, unknown>>;
+
+async function main() {
+  // ── Level 1: Auto-billing via biller.ai() wrapper ──
+
+  const aii = biller.ai(ai, {
+    inputDebit: { tag: "PREMIUM_CALL" },
+    outputDebit: { tag: "EXTRA_FEE" },
+  });
+
+  console.log("--- Level 1: biller.ai() auto-wrapper (streamText) ---");
+
+  const wrappedStreamText = aii.streamText as WrappedFn;
+  const result = await wrappedStreamText({
+    userId: "c0971bcb-b901-4c3e-a191-c9a97871c39f",
+    model: openai("gpt-4o-mini"),
+    prompt: "Write a 2 sentence story about a robot.",
+  });
+  console.log(
+    `  Generated: "${((await result.text) as string).slice(0, 80)}..."\n`
+  );
+
+  // ── Level 1: With user's own onStepFinish ──
+
+  console.log("--- Level 1: With user onStepFinish callback ---");
+
+  await wrappedStreamText({
+    userId: "c0971bcb-b901-4c3e-a191-c9a97871c39f",
+    model: openai("gpt-4o-mini"),
+    prompt: "Say hello in Spanish.",
+    onStepFinish: (event: {
+      stepNumber: number;
+      usage: { totalTokens?: number };
+    }) => {
+      console.log(
+        `  Step ${event.stepNumber}: ${event.usage.totalTokens ?? 0} tokens`
+      );
+    },
+  });
+
+  console.log();
+
+  // ── Level 2: Manual biller.trackAI() ──
+
+  console.log("--- Level 2: Manual biller.trackAI() in onStepFinish ---");
+
+  const manualResult = await ai.streamText({
+    model: openai("gpt-4o-mini"),
+    prompt: "Say hello in French.",
+    onStepFinish: (event: {
+      model: { modelId: string; provider: string };
+      usage: {
+        inputTokens?: number;
+        outputTokens?: number;
+        totalTokens?: number;
+      };
+    }) => {
+      biller.trackAI(
+        "c0971bcb-b901-4c3e-a191-c9a97871c39f",
+        { modelId: event.model.modelId, provider: event.model.provider },
+        {
+          inputTokens: event.usage.inputTokens ?? 0,
+          outputTokens: event.usage.outputTokens ?? 0,
+          totalTokens:
+            (event.usage.inputTokens ?? 0) + (event.usage.outputTokens ?? 0),
+        },
+        { userId: "" },
+        {
+          inputDebit: { tag: "PREMIUM_CALL" },
+          outputDebit: { tag: "EXTRA_FEE" },
+          inputCacheDebit: { tag: "PREMIUM_CALL" },
+          outputCacheDebit: { tag: "EXTRA_FEE" },
+        }
+      );
+      console.log(`  Tracked ${event.usage.totalTokens} tokens`);
+    },
+  });
+
+  console.log(`  Generated: "${await manualResult.text}"\n`);
+  console.log("All AI SDK wrapper examples completed.");
+}
+
+main().catch(console.error);
diff --git a/examples/bun.lock b/examples/bun.lock
index 0f31057..4de691a 100644
--- a/examples/bun.lock
+++ b/examples/bun.lock
@@ -5,7 +5,10 @@
     "": {
       "name": "@scrawn/sdkexamples",
       "dependencies": {
+        "@ai-sdk/openai": "^3.0.65",
+        "@scrawn/analytics": "link:@scrawn/analytics",
         "@scrawn/core": "link:@scrawn/core",
+        "ai": "^6.0.191",
         "dotenv": "^17.2.3",
         "express": "^4.18.2",
       },
@@ -16,8 +19,22 @@
     },
   },
   "packages": {
+    "@ai-sdk/gateway": ["@ai-sdk/gateway@3.0.120", "", { "dependencies": { "@ai-sdk/provider": "3.0.10", "@ai-sdk/provider-utils": "4.0.27", "@vercel/oidc": "3.2.0" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-MYKAeD2q7/sa1ZdqtL2tw0Me0B8Tok6Q/fhkJDhJl39dG8u+VBlWO9yk9lcdm784bM418o1EKObo4aOxs6+18Q=="],
+
+    "@ai-sdk/openai": ["@ai-sdk/openai@3.0.65", "", { "dependencies": { "@ai-sdk/provider": "3.0.10", "@ai-sdk/provider-utils": "4.0.27" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-ZlVoWH+zrdiYDiUt6n/xvfCsk33mzsB81TUQkBRVx79rxU1FKZqVH9J/QCtEpSLqx0cUzjvtIw9l9p7EbUv+dw=="],
+
+    "@ai-sdk/provider": ["@ai-sdk/provider@3.0.10", "", { "dependencies": { "json-schema": "^0.4.0" } }, "sha512-Q3BZ27qfpYqnCYGvE3vt+Qi6LGOF9R5Nmzn+9JoM1lCRsD9mYaIhfJLkSunN48nfGXJ6n+XNV0J/XVpqGQl7Dw=="],
+
+    "@ai-sdk/provider-utils": ["@ai-sdk/provider-utils@4.0.27", "", { "dependencies": { "@ai-sdk/provider": "3.0.10", "@standard-schema/spec": "^1.1.0", "eventsource-parser": "^3.0.8" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-ubkAJ+xODouwtmN1tYlvTPphH1hPOBfZaEQe8U7skGvFAnIRs9PPpsq57bC2+Ky/MB4yzhd6YOsxTAx9sGpazw=="],
+
+    "@opentelemetry/api": ["@opentelemetry/api@1.9.1", "", {}, "sha512-gLyJlPHPZYdAk1JENA9LeHejZe1Ti77/pTeFm/nMXmQH/HFZlcS/O2XJB+L8fkbrNSqhdtlvjBVjxwUYanNH5Q=="],
+
+    "@scrawn/analytics": ["@scrawn/analytics@link:@scrawn/analytics", {}],
+
     "@scrawn/core": ["@scrawn/core@link:@scrawn/core", {}],
 
+    "@standard-schema/spec": ["@standard-schema/spec@1.1.0", "", {}, "sha512-l2aFy5jALhniG5HgqrD6jXLi/rUWrKvqN/qJx6yoJsgKhblVd+iqqU4RCXavm/jPityDo5TCvKMnpjKnOriy0w=="],
+
     "@types/body-parser": ["@types/body-parser@1.19.6", "", { "dependencies": { "@types/connect": "*", "@types/node": "*" } }, "sha512-HLFeCYgz89uk22N5Qg3dvGvsv46B8GLvKKo1zKG4NybA8U2DiEO3w9lqGg29t/tfLRJpJ6iQxnVw4OnB7MoM9g=="],
 
     "@types/connect": ["@types/connect@3.4.38", "", { "dependencies": { "@types/node": "*" } }, "sha512-K6uROf1LD88uDQqJCktA4yzL1YYAK6NgfsI0v/mTgyPKWsX1CnJ0XPSDhViejru1GcRkLWb8RlzFYJRqGUbaug=="],
@@ -40,8 +57,12 @@
 
     "@types/serve-static": ["@types/serve-static@1.15.10", "", { "dependencies": { "@types/http-errors": "*", "@types/node": "*", "@types/send": "<1" } }, "sha512-tRs1dB+g8Itk72rlSI2ZrW6vZg0YrLI81iQSTkMmOqnqCaNr/8Ek4VwWcN5vZgCYWbg/JJSGBlUaYGAOP73qBw=="],
 
+    "@vercel/oidc": ["@vercel/oidc@3.2.0", "", {}, "sha512-UycprH3T6n3jH0k44NHMa7pnFHGu/N05MjojYr+Mc6I7obkoLIJujSWwin1pCvdy/eOxrI/l3uDLQsmcrOb4ug=="],
+
     "accepts": ["accepts@1.3.8", "", { "dependencies": { "mime-types": "~2.1.34", "negotiator": "0.6.3" } }, "sha512-PYAthTa2m2VKxuvSD3DPC/Gy+U+sOA1LAuT8mkmRuvw+NACSaeXEQ+NHcVF7rONl6qcaxV3Uuemwawk+7+SJLw=="],
 
+    "ai": ["ai@6.0.191", "", { "dependencies": { "@ai-sdk/gateway": "3.0.120", "@ai-sdk/provider": "3.0.10", "@ai-sdk/provider-utils": "4.0.27", "@opentelemetry/api": "^1.9.0" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-zAxvjKebQE7YkSyyNIl0OM7i6/zygnKeF+yNUjD4nWOelYrG+LpDd6RnH6mjySI4zUpZ7o4wbnmAy8jc6u98vQ=="],
+
     "array-flatten": ["array-flatten@1.1.1", "", {}, "sha512-PCVAQswWemu6UdxsDFFX/+gVeYqKAod3D3UVm91jHwynguOwAvYPhx8nNlM++NqRcK6CxxpUafjmhIdKiHibqg=="],
 
     "body-parser": ["body-parser@1.20.3", "", { "dependencies": { "bytes": "3.1.2", "content-type": "~1.0.5", "debug": "2.6.9", "depd": "2.0.0", "destroy": "1.2.0", "http-errors": "2.0.0", "iconv-lite": "0.4.24", "on-finished": "2.4.1", "qs": "6.13.0", "raw-body": "2.5.2", "type-is": "~1.6.18", "unpipe": "1.0.0" } }, "sha512-7rAxByjUMqQ3/bHJy7D6OGXvx/MMc4IqBn/X0fcM1QUcAItpZrBEYhWGem+tzXH90c+G01ypMcYJBO9Y30203g=="],
@@ -84,6 +105,8 @@
 
     "etag": ["etag@1.8.1", "", {}, "sha512-aIL5Fx7mawVa300al2BnEE4iNvo1qETxLrPI/o05L7z6go7fCw1J6EQmbK4FmJ2AS7kgVF/KEZWufBfdClMcPg=="],
 
+    "eventsource-parser": ["eventsource-parser@3.0.8", "", {}, "sha512-70QWGkr4snxr0OXLRWsFLeRBIRPuQOvt4s8QYjmUlmlkyTZkRqS7EDVRZtzU3TiyDbXSzaOeF0XUKy8PchzukQ=="],
+
     "express": ["express@4.21.2", "", { "dependencies": { "accepts": "~1.3.8", "array-flatten": "1.1.1", "body-parser": "1.20.3", "content-disposition": "0.5.4", "content-type": "~1.0.4", "cookie": "0.7.1", "cookie-signature": "1.0.6", "debug": "2.6.9", "depd": "2.0.0", "encodeurl": "~2.0.0", "escape-html": "~1.0.3", "etag": "~1.8.1", "finalhandler": "1.3.1", "fresh": "0.5.2", "http-errors": "2.0.0", "merge-descriptors": "1.0.3", "methods": "~1.1.2", "on-finished": "2.4.1", "parseurl": "~1.3.3", "path-to-regexp": "0.1.12", "proxy-addr": "~2.0.7", "qs": "6.13.0", "range-parser": "~1.2.1", "safe-buffer": "5.2.1", "send": "0.19.0", "serve-static": "1.16.2", "setprototypeof": "1.2.0", "statuses": "2.0.1", "type-is": "~1.6.18", "utils-merge": "1.0.1", "vary": "~1.1.2" } }, "sha512-28HqgMZAmih1Czt9ny7qr6ek2qddF4FclbMzwhCREB6OFfH+rXAnuNCwo1/wFvrtbgsQDb4kSbX9de9lFbrXnA=="],
 
     "finalhandler": ["finalhandler@1.3.1", "", { "dependencies": { "debug": "2.6.9", "encodeurl": "~2.0.0", "escape-html": "~1.0.3", "on-finished": "2.4.1", "parseurl": "~1.3.3", "statuses": "2.0.1", "unpipe": "~1.0.0" } }, "sha512-6BN9trH7bp3qvnrRyzsBz+g3lZxTNZTbVO2EV1CS0WIcDbawYVdYvGflME/9QP0h0pYlCDBCTjYa9nZzMDpyxQ=="],
@@ -112,6 +135,8 @@
 
     "ipaddr.js": ["ipaddr.js@1.9.1", "", {}, "sha512-0KI/607xoxSToH7GjN1FfSbLoU0+btTicjsQSWQlh/hZykN8KpmMf7uYwPW3R+akZ6R/w18ZlXSHBYXiYUPO3g=="],
 
+    "json-schema": ["json-schema@0.4.0", "", {}, "sha512-es94M3nTIfsEPisRafak+HDLfHXnKBhV3vU5eqPcS3flIWqcxJWgXHXiey3YrpaNsanY5ei1VoYEbOzijuq9BA=="],
+
     "math-intrinsics": ["math-intrinsics@1.1.0", "", {}, "sha512-/IXtbwEk5HTPyEwyKX6hGkYXxM9nbj64B+ilVJnC/R6B0pH5G4V3b0pVbL7DBj4tkhBAppbQUlf6F6Xl9LHu1g=="],
 
     "media-typer": ["media-typer@0.3.0", "", {}, "sha512-dq+qelQ9akHpcOl/gUVRTxVIOkAJ1wR3QAvb4RsVjS8oVoFjDGTc679wJYmUmknUF5HwMLOgb5O+a3KxfWapPQ=="],
@@ -178,6 +203,8 @@
 
     "vary": ["vary@1.1.2", "", {}, "sha512-BNGbWLfd0eUPabhkXUVm0j8uuvREyTh5ovRa/dyow/BqAbZJyC+5fU+IzQOzmAKzYqYRAISoRhdQr3eIZ/PXqg=="],
 
+    "zod": ["zod@4.4.3", "", {}, "sha512-ytENFjIJFl2UwYglde2jchW2Hwm4GJFLDiSXWdTrJQBIN9Fcyp7n4DhxJEiWNAJMV1/BqWfW/kkg71UDcHJyTQ=="],
+
     "@types/serve-static/@types/send": ["@types/send@0.17.6", "", { "dependencies": { "@types/mime": "^1", "@types/node": "*" } }, "sha512-Uqt8rPBE8SY0RK8JB1EzVOIZ32uqy8HwdxCnoCOsYrvnswqmFZ/k+9Ikidlk/ImhsdvBsloHbAlewb2IEBV/Og=="],
 
     "send/encodeurl": ["encodeurl@1.0.2", "", {}, "sha512-TPJXq8JqFaVYm2CWmPvnP2Iyo4ZSM7/QKcSmuMLDObfpH5fi7RUGmd/rTDf+rut/saiDiQEeVTNgAmJEdAOx0w=="],
diff --git a/examples/package.json b/examples/package.json
index 059eaa8..0f924c6 100644
--- a/examples/package.json
+++ b/examples/package.json
@@ -13,9 +13,11 @@
     "@types/node": "^24.10.0"
   },
   "dependencies": {
-    "dotenv": "^17.2.3",
-    "express": "^4.18.2",
+    "@ai-sdk/openai": "^3.0.65",
+    "@scrawn/analytics": "link:@scrawn/analytics",
     "@scrawn/core": "link:@scrawn/core",
-    "@scrawn/analytics": "link:@scrawn/analytics"
+    "ai": "^6.0.191",
+    "dotenv": "^17.2.3",
+    "express": "^4.18.2"
   }
 }

From 4299f7d0bc6b9da6f428cd4a7f2c8b7ee57b4d15 Mon Sep 17 00:00:00 2001
From: Devyash Saini <dysaini2004@gmail.com>
Date: Tue, 26 May 2026 22:18:27 +0530
Subject: [PATCH 04/11] fix(ai): properly type biller.ai() return to eliminate
 user-side casts

- Return Record<string, (params) => Promise<unknown>> instead of Record<string, unknown>
- Example now works cleanly without any as casts
---
 examples/ai-sdk-wrapper-usage.ts   | 35 +++++++++++++-----------------
 packages/scrawn/src/core/scrawn.ts | 21 +++++++++---------
 2 files changed, 26 insertions(+), 30 deletions(-)

diff --git a/examples/ai-sdk-wrapper-usage.ts b/examples/ai-sdk-wrapper-usage.ts
index 223a3d2..a4e4b95 100644
--- a/examples/ai-sdk-wrapper-usage.ts
+++ b/examples/ai-sdk-wrapper-usage.ts
@@ -4,11 +4,6 @@ import { biller } from "./scrawn/biller.js";
 import { config as dotenvConfig } from "dotenv";
 dotenvConfig({ path: ".env.local" });
 
-// Helper: cast the wrapped function back to a usable callable type
-type WrappedFn = (
-  params: Record<string, unknown>
-) => Promise<Record<string, unknown>>;
-
 async function main() {
   // ── Level 1: Auto-billing via biller.ai() wrapper ──
 
@@ -19,35 +14,35 @@ async function main() {
 
   console.log("--- Level 1: biller.ai() auto-wrapper (streamText) ---");
 
-  const wrappedStreamText = aii.streamText as WrappedFn;
-  const result = await wrappedStreamText({
+  const result = await aii.streamText({
     userId: "c0971bcb-b901-4c3e-a191-c9a97871c39f",
     model: openai("gpt-4o-mini"),
     prompt: "Write a 2 sentence story about a robot.",
+    onStepFinish: (event: {
+      stepNumber: number;
+      usage: { totalTokens?: number };
+    }) => {
+      console.log(
+        `  Step ${event.stepNumber}: ${event.usage.totalTokens ?? 0} tokens`
+      );
+    },
   });
+
   console.log(
     `  Generated: "${((await result.text) as string).slice(0, 80)}..."\n`
   );
 
-  // ── Level 1: With user's own onStepFinish ──
+  // ── Level 1: generateText (non-streaming) ──
 
-  console.log("--- Level 1: With user onStepFinish callback ---");
+  console.log("--- Level 1: biller.ai() auto-wrapper (generateText) ---");
 
-  await wrappedStreamText({
+  const genResult = await aii.generateText({
     userId: "c0971bcb-b901-4c3e-a191-c9a97871c39f",
     model: openai("gpt-4o-mini"),
-    prompt: "Say hello in Spanish.",
-    onStepFinish: (event: {
-      stepNumber: number;
-      usage: { totalTokens?: number };
-    }) => {
-      console.log(
-        `  Step ${event.stepNumber}: ${event.usage.totalTokens ?? 0} tokens`
-      );
-    },
+    prompt: "What is 2+2?",
   });
 
-  console.log();
+  console.log(`  Answer: ${(genResult as { text: string }).text}\n`);
 
   // ── Level 2: Manual biller.trackAI() ──
 
diff --git a/packages/scrawn/src/core/scrawn.ts b/packages/scrawn/src/core/scrawn.ts
index cb5c425..175ef6b 100644
--- a/packages/scrawn/src/core/scrawn.ts
+++ b/packages/scrawn/src/core/scrawn.ts
@@ -1193,39 +1193,40 @@ export class Scrawn<
   /**
    * Wraps the Vercel AI SDK with automatic per-step billing.
    *
-   * Returns a proxied version of the AI SDK where `streamText`, `generateText`,
-   * `streamObject`, and `generateObject` automatically:
-   * 1. Accept a `userId` parameter for billing
-   * 2. Fire-and-forget `aiTokenStreamConsumer` on every step via `onStepFinish`
-   * 3. Chain the user's own `onStepFinish`/`onFinish` callbacks
+   * Returns a proxied version where `streamText`, `generateText`,
+   * `streamObject`, and `generateObject` accept a `userId` parameter and
+   * automatically track token usage via `onStepFinish`.
+   *
+   * User callbacks (`onStepFinish`, `onFinish`) are chained alongside billing.
    *
    * @param sdk - The Vercel AI SDK module (import * as ai from "ai")
    * @param opts - Default billing configuration for all calls
-   * @returns A proxied AI SDK with billing injected
    *
    * @example
    * ```typescript
    * import * as ai from "ai";
-   * import { biller } from "./scrawn/biller";
    *
    * const aii = biller.ai(ai, {
    *   inputDebit: { tag: "AI_INPUT" },
    *   outputDebit: { tag: "AI_OUTPUT" },
    * });
    *
+   * // No casts — just add userId
    * const result = await aii.streamText({
    *   userId: "user-123",
    *   model: openai("gpt-4o-mini"),
    *   prompt: "Write a story.",
-   *   onStepFinish: event => { console.log("Step done"); },
    * });
    * ```
    */
   ai(
     sdk: Record<string, unknown>,
     opts: BillableAIOptions<TTags>
-  ): Record<string, unknown> {
-    return createBillableAI(sdk, this, opts);
+  ): Record<string, (params: Record<string, unknown>) => Promise<unknown>> {
+    return createBillableAI(sdk, this, opts) as Record<
+      string,
+      (params: Record<string, unknown>) => Promise<unknown>
+    >;
   }
 
   /**

From 211d8284c1066e1bf5ee81635c0143b03b2ee1d0 Mon Sep 17 00:00:00 2001
From: Devyash Saini <dysaini2004@gmail.com>
Date: Tue, 26 May 2026 22:24:06 +0530
Subject: [PATCH 05/11] fix(ai): use const TSDK intersection to preserve all AI
 SDK return types
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- biller.ai() now infers TSDK from the input module via const generic
- Returns TSDK & overlay — original types (including return) survive
- streamText still returns StreamTextResult, generateText returns GenerateTextResult
- Example now has zero casts: no as WrappedFn, no as string, no as unknown
- Removed unnecessary BillableAI/StreamTextResult/GenerateTextResult types
---
 examples/ai-sdk-wrapper-usage.ts     |  8 +++-----
 packages/scrawn/src/core/ai/types.ts | 11 +++++++++++
 packages/scrawn/src/core/scrawn.ts   | 22 +++++++++++-----------
 3 files changed, 25 insertions(+), 16 deletions(-)

diff --git a/examples/ai-sdk-wrapper-usage.ts b/examples/ai-sdk-wrapper-usage.ts
index a4e4b95..bd2cb1d 100644
--- a/examples/ai-sdk-wrapper-usage.ts
+++ b/examples/ai-sdk-wrapper-usage.ts
@@ -28,9 +28,7 @@ async function main() {
     },
   });
 
-  console.log(
-    `  Generated: "${((await result.text) as string).slice(0, 80)}..."\n`
-  );
+  console.log(`  Generated: "${(await result.text).slice(0, 80)}..."\n`);
 
   // ── Level 1: generateText (non-streaming) ──
 
@@ -42,11 +40,11 @@ async function main() {
     prompt: "What is 2+2?",
   });
 
-  console.log(`  Answer: ${(genResult as { text: string }).text}\n`);
+  console.log(`  Answer: ${genResult.text}\n`);
 
   // ── Level 2: Manual biller.trackAI() ──
 
-  console.log("--- Level 2: Manual biller.trackAI() in onStepFinish ---");
+  console.log("--- Level 2: Manual biller.trackAI() ---");
 
   const manualResult = await ai.streamText({
     model: openai("gpt-4o-mini"),
diff --git a/packages/scrawn/src/core/ai/types.ts b/packages/scrawn/src/core/ai/types.ts
index 6323abc..5e324e7 100644
--- a/packages/scrawn/src/core/ai/types.ts
+++ b/packages/scrawn/src/core/ai/types.ts
@@ -69,3 +69,14 @@ export interface ModelInfo {
   /** Provider name, e.g. "openai", "anthropic". */
   provider: string;
 }
+
+/**
+ * Minimal subset of the AI SDK model info needed for billing.
+ * Comes from OnStepFinishEvent.model or OnFinishEvent.model.
+ */
+export interface ModelInfo {
+  /** Model ID, e.g. "gpt-4o-mini". */
+  modelId: string;
+  /** Provider name, e.g. "openai", "anthropic". */
+  provider: string;
+}
diff --git a/packages/scrawn/src/core/scrawn.ts b/packages/scrawn/src/core/scrawn.ts
index 175ef6b..e9e2e7c 100644
--- a/packages/scrawn/src/core/scrawn.ts
+++ b/packages/scrawn/src/core/scrawn.ts
@@ -1193,9 +1193,11 @@ export class Scrawn<
   /**
    * Wraps the Vercel AI SDK with automatic per-step billing.
    *
-   * Returns a proxied version where `streamText`, `generateText`,
-   * `streamObject`, and `generateObject` accept a `userId` parameter and
-   * automatically track token usage via `onStepFinish`.
+   * Returns the AI SDK with `streamText`, `generateText`, `streamObject`,
+   * and `generateObject` patched to accept a `userId` parameter and
+   * automatically track token usage. All original AI SDK types pass
+   * through unchanged — returns the same module shape you passed in,
+   * with billing injected.
    *
    * User callbacks (`onStepFinish`, `onFinish`) are chained alongside billing.
    *
@@ -1211,22 +1213,20 @@ export class Scrawn<
    *   outputDebit: { tag: "AI_OUTPUT" },
    * });
    *
-   * // No casts — just add userId
    * const result = await aii.streamText({
    *   userId: "user-123",
    *   model: openai("gpt-4o-mini"),
    *   prompt: "Write a story.",
    * });
+   * // result.text → Promise<string> (preserved from AI SDK)
    * ```
    */
-  ai(
-    sdk: Record<string, unknown>,
+  ai<const TSDK extends Record<string, unknown>>(
+    sdk: TSDK,
     opts: BillableAIOptions<TTags>
-  ): Record<string, (params: Record<string, unknown>) => Promise<unknown>> {
-    return createBillableAI(sdk, this, opts) as Record<
-      string,
-      (params: Record<string, unknown>) => Promise<unknown>
-    >;
+  ): TSDK & Record<string, (params: Record<string, unknown>) => unknown> {
+    return createBillableAI(sdk, this, opts) as TSDK &
+      Record<string, (params: Record<string, unknown>) => unknown>;
   }
 
   /**

From d6d0a0f6b2888df43f2059912b8b8a6b7bf75da9 Mon Sep 17 00:00:00 2001
From: Devyash Saini <dysaini2004@gmail.com>
Date: Tue, 26 May 2026 22:26:06 +0530
Subject: [PATCH 06/11] fix(ai): use Omit to replace wrapped fns with widened
 userId params

- Omit streamText/generateText/streamObject/generateObject from TSDK
- Re-add them with Record<string, unknown> params accepting userId
- Return types preserved: streamText returns {text: Promise<string>}, generateText returns {text: string}
---
 packages/scrawn/src/core/scrawn.ts | 29 ++++++++++++++++++++++++++---
 1 file changed, 26 insertions(+), 3 deletions(-)

diff --git a/packages/scrawn/src/core/scrawn.ts b/packages/scrawn/src/core/scrawn.ts
index e9e2e7c..7393a09 100644
--- a/packages/scrawn/src/core/scrawn.ts
+++ b/packages/scrawn/src/core/scrawn.ts
@@ -1224,9 +1224,32 @@ export class Scrawn<
   ai<const TSDK extends Record<string, unknown>>(
     sdk: TSDK,
     opts: BillableAIOptions<TTags>
-  ): TSDK & Record<string, (params: Record<string, unknown>) => unknown> {
-    return createBillableAI(sdk, this, opts) as TSDK &
-      Record<string, (params: Record<string, unknown>) => unknown>;
+  ): Omit<
+    TSDK,
+    "streamText" | "generateText" | "streamObject" | "generateObject"
+  > & {
+    streamText: (
+      params: Record<string, unknown>
+    ) => Promise<{ text: Promise<string> } & Record<string, unknown>>;
+    generateText: (
+      params: Record<string, unknown>
+    ) => Promise<{ text: string } & Record<string, unknown>>;
+    streamObject: (params: Record<string, unknown>) => Promise<unknown>;
+    generateObject: (params: Record<string, unknown>) => Promise<unknown>;
+  } {
+    return createBillableAI(sdk, this, opts) as Omit<
+      TSDK,
+      "streamText" | "generateText" | "streamObject" | "generateObject"
+    > & {
+      streamText: (
+        params: Record<string, unknown>
+      ) => Promise<{ text: Promise<string> } & Record<string, unknown>>;
+      generateText: (
+        params: Record<string, unknown>
+      ) => Promise<{ text: string } & Record<string, unknown>>;
+      streamObject: (params: Record<string, unknown>) => Promise<unknown>;
+      generateObject: (params: Record<string, unknown>) => Promise<unknown>;
+    };
   }
 
   /**

From 84db5683797703242c51b30da544a67e70d8918b Mon Sep 17 00:00:00 2001
From: Devyash Saini <dysaini2004@gmail.com>
Date: Tue, 26 May 2026 22:31:58 +0530
Subject: [PATCH 07/11] fix(ai): use conditional inference (infer P, infer R)
 to preserve exact AI SDK types

- WithUserId<TSDK> type: Omit wrapped fns, re-add with P & BillableCallParams
- Preserves original AI SDK param types (model, prompt, tools, etc.)
- Preserves original return types (StreamTextResult, GenerateTextResult)
- User gets full AI SDK intellisense + userId
---
 packages/scrawn/src/core/ai/types.ts | 40 ++++++++++++++++++++++------
 packages/scrawn/src/core/scrawn.ts   | 29 +++-----------------
 packages/scrawn/src/index.ts         |  1 +
 3 files changed, 36 insertions(+), 34 deletions(-)

diff --git a/packages/scrawn/src/core/ai/types.ts b/packages/scrawn/src/core/ai/types.ts
index 5e324e7..a80ef93 100644
--- a/packages/scrawn/src/core/ai/types.ts
+++ b/packages/scrawn/src/core/ai/types.ts
@@ -70,13 +70,37 @@ export interface ModelInfo {
   provider: string;
 }
 
+// ── Billable AI SDK type mapping ──
+
+/** Keys in the AI SDK that we wrap with billing. */
+type BillableKeys =
+  | "streamText"
+  | "generateText"
+  | "streamObject"
+  | "generateObject";
+
 /**
- * Minimal subset of the AI SDK model info needed for billing.
- * Comes from OnStepFinishEvent.model or OnFinishEvent.model.
+ * Given the original AI SDK type TSDK, replaces the wrapped function
+ * signatures with widened versions that accept `userId`. Preserves the
+ * original param types via conditional inference and the original return
+ * types unchanged.
+ *
+ * All other AI SDK properties pass through via Omit.
  */
-export interface ModelInfo {
-  /** Model ID, e.g. "gpt-4o-mini". */
-  modelId: string;
-  /** Provider name, e.g. "openai", "anthropic". */
-  provider: string;
-}
+export type WithUserId<TSDK extends Record<string, unknown>> = Omit<
+  TSDK,
+  BillableKeys
+> & {
+  streamText: TSDK extends { streamText: (params: infer P) => infer R }
+    ? (params: P & BillableCallParams) => R
+    : never;
+  generateText: TSDK extends { generateText: (params: infer P) => infer R }
+    ? (params: P & BillableCallParams) => R
+    : never;
+  streamObject: TSDK extends { streamObject: (params: infer P) => infer R }
+    ? (params: P & BillableCallParams) => R
+    : never;
+  generateObject: TSDK extends { generateObject: (params: infer P) => infer R }
+    ? (params: P & BillableCallParams) => R
+    : never;
+};
diff --git a/packages/scrawn/src/core/scrawn.ts b/packages/scrawn/src/core/scrawn.ts
index 7393a09..e008fc2 100644
--- a/packages/scrawn/src/core/scrawn.ts
+++ b/packages/scrawn/src/core/scrawn.ts
@@ -62,6 +62,7 @@ import {
 } from "./pricing/index.js";
 import { createBillableAI } from "./ai/wrap.js";
 import type { BillableAIOptions } from "./ai/types.js";
+import type { WithUserId } from "./ai/types.js";
 import { buildAIPayload } from "./ai/track.js";
 import type {
   LanguageModelUsage,
@@ -1224,32 +1225,8 @@ export class Scrawn<
   ai<const TSDK extends Record<string, unknown>>(
     sdk: TSDK,
     opts: BillableAIOptions<TTags>
-  ): Omit<
-    TSDK,
-    "streamText" | "generateText" | "streamObject" | "generateObject"
-  > & {
-    streamText: (
-      params: Record<string, unknown>
-    ) => Promise<{ text: Promise<string> } & Record<string, unknown>>;
-    generateText: (
-      params: Record<string, unknown>
-    ) => Promise<{ text: string } & Record<string, unknown>>;
-    streamObject: (params: Record<string, unknown>) => Promise<unknown>;
-    generateObject: (params: Record<string, unknown>) => Promise<unknown>;
-  } {
-    return createBillableAI(sdk, this, opts) as Omit<
-      TSDK,
-      "streamText" | "generateText" | "streamObject" | "generateObject"
-    > & {
-      streamText: (
-        params: Record<string, unknown>
-      ) => Promise<{ text: Promise<string> } & Record<string, unknown>>;
-      generateText: (
-        params: Record<string, unknown>
-      ) => Promise<{ text: string } & Record<string, unknown>>;
-      streamObject: (params: Record<string, unknown>) => Promise<unknown>;
-      generateObject: (params: Record<string, unknown>) => Promise<unknown>;
-    };
+  ): WithUserId<TSDK> {
+    return createBillableAI(sdk, this, opts) as WithUserId<TSDK>;
   }
 
   /**
diff --git a/packages/scrawn/src/index.ts b/packages/scrawn/src/index.ts
index 4eb769d..f938db8 100644
--- a/packages/scrawn/src/index.ts
+++ b/packages/scrawn/src/index.ts
@@ -68,4 +68,5 @@ export type {
   BillableCallParams,
   LanguageModelUsage,
   ModelInfo,
+  WithUserId,
 } from "./core/ai/types.ts";

From a5be887c88e4464c2c8e714f9ef8fe81c257970a Mon Sep 17 00:00:00 2001
From: Devyash Saini <dysaini2004@gmail.com>
Date: Tue, 26 May 2026 23:24:42 +0530
Subject: [PATCH 08/11] fix: update ai sdk wrapper example

Signed-off-by: Devyash Saini <dysaini2004@gmail.com>
---
 examples/ai-sdk-wrapper-usage.ts | 73 +++-----------------------------
 examples/bun.lock                |  3 ++
 examples/package.json            |  1 +
 3 files changed, 9 insertions(+), 68 deletions(-)

diff --git a/examples/ai-sdk-wrapper-usage.ts b/examples/ai-sdk-wrapper-usage.ts
index bd2cb1d..cd00a12 100644
--- a/examples/ai-sdk-wrapper-usage.ts
+++ b/examples/ai-sdk-wrapper-usage.ts
@@ -1,85 +1,22 @@
 import * as ai from "ai";
-import { openai } from "@ai-sdk/openai";
+import { google } from "@ai-sdk/google";
 import { biller } from "./scrawn/biller.js";
-import { config as dotenvConfig } from "dotenv";
-dotenvConfig({ path: ".env.local" });
+import { config } from "dotenv";
+config({ path: ".env.local" });
 
 async function main() {
-  // ── Level 1: Auto-billing via biller.ai() wrapper ──
-
   const aii = biller.ai(ai, {
     inputDebit: { tag: "PREMIUM_CALL" },
     outputDebit: { tag: "EXTRA_FEE" },
   });
 
-  console.log("--- Level 1: biller.ai() auto-wrapper (streamText) ---");
-
   const result = await aii.streamText({
     userId: "c0971bcb-b901-4c3e-a191-c9a97871c39f",
-    model: openai("gpt-4o-mini"),
+    model: google("gemini-2.5-flash"),
     prompt: "Write a 2 sentence story about a robot.",
-    onStepFinish: (event: {
-      stepNumber: number;
-      usage: { totalTokens?: number };
-    }) => {
-      console.log(
-        `  Step ${event.stepNumber}: ${event.usage.totalTokens ?? 0} tokens`
-      );
-    },
-  });
-
-  console.log(`  Generated: "${(await result.text).slice(0, 80)}..."\n`);
-
-  // ── Level 1: generateText (non-streaming) ──
-
-  console.log("--- Level 1: biller.ai() auto-wrapper (generateText) ---");
-
-  const genResult = await aii.generateText({
-    userId: "c0971bcb-b901-4c3e-a191-c9a97871c39f",
-    model: openai("gpt-4o-mini"),
-    prompt: "What is 2+2?",
-  });
-
-  console.log(`  Answer: ${genResult.text}\n`);
-
-  // ── Level 2: Manual biller.trackAI() ──
-
-  console.log("--- Level 2: Manual biller.trackAI() ---");
-
-  const manualResult = await ai.streamText({
-    model: openai("gpt-4o-mini"),
-    prompt: "Say hello in French.",
-    onStepFinish: (event: {
-      model: { modelId: string; provider: string };
-      usage: {
-        inputTokens?: number;
-        outputTokens?: number;
-        totalTokens?: number;
-      };
-    }) => {
-      biller.trackAI(
-        "c0971bcb-b901-4c3e-a191-c9a97871c39f",
-        { modelId: event.model.modelId, provider: event.model.provider },
-        {
-          inputTokens: event.usage.inputTokens ?? 0,
-          outputTokens: event.usage.outputTokens ?? 0,
-          totalTokens:
-            (event.usage.inputTokens ?? 0) + (event.usage.outputTokens ?? 0),
-        },
-        { userId: "" },
-        {
-          inputDebit: { tag: "PREMIUM_CALL" },
-          outputDebit: { tag: "EXTRA_FEE" },
-          inputCacheDebit: { tag: "PREMIUM_CALL" },
-          outputCacheDebit: { tag: "EXTRA_FEE" },
-        }
-      );
-      console.log(`  Tracked ${event.usage.totalTokens} tokens`);
-    },
   });
 
-  console.log(`  Generated: "${await manualResult.text}"\n`);
-  console.log("All AI SDK wrapper examples completed.");
+  console.log(`Generated: "${await result.text}"\n`);
 }
 
 main().catch(console.error);
diff --git a/examples/bun.lock b/examples/bun.lock
index 4de691a..0cb74e5 100644
--- a/examples/bun.lock
+++ b/examples/bun.lock
@@ -5,6 +5,7 @@
     "": {
       "name": "@scrawn/sdkexamples",
       "dependencies": {
+        "@ai-sdk/google": "^3.0.79",
         "@ai-sdk/openai": "^3.0.65",
         "@scrawn/analytics": "link:@scrawn/analytics",
         "@scrawn/core": "link:@scrawn/core",
@@ -21,6 +22,8 @@
   "packages": {
     "@ai-sdk/gateway": ["@ai-sdk/gateway@3.0.120", "", { "dependencies": { "@ai-sdk/provider": "3.0.10", "@ai-sdk/provider-utils": "4.0.27", "@vercel/oidc": "3.2.0" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-MYKAeD2q7/sa1ZdqtL2tw0Me0B8Tok6Q/fhkJDhJl39dG8u+VBlWO9yk9lcdm784bM418o1EKObo4aOxs6+18Q=="],
 
+    "@ai-sdk/google": ["@ai-sdk/google@3.0.79", "", { "dependencies": { "@ai-sdk/provider": "3.0.10", "@ai-sdk/provider-utils": "4.0.27" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-QWVAvYeA7JzEX2wkSyXOWv/I9PD9kvTzdykkSTLi+Eu8RyJ6gA0tdPIGa8esEtOcHE//G5vy6FTB70qQw8l/uw=="],
+
     "@ai-sdk/openai": ["@ai-sdk/openai@3.0.65", "", { "dependencies": { "@ai-sdk/provider": "3.0.10", "@ai-sdk/provider-utils": "4.0.27" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-ZlVoWH+zrdiYDiUt6n/xvfCsk33mzsB81TUQkBRVx79rxU1FKZqVH9J/QCtEpSLqx0cUzjvtIw9l9p7EbUv+dw=="],
 
     "@ai-sdk/provider": ["@ai-sdk/provider@3.0.10", "", { "dependencies": { "json-schema": "^0.4.0" } }, "sha512-Q3BZ27qfpYqnCYGvE3vt+Qi6LGOF9R5Nmzn+9JoM1lCRsD9mYaIhfJLkSunN48nfGXJ6n+XNV0J/XVpqGQl7Dw=="],
diff --git a/examples/package.json b/examples/package.json
index 0f924c6..2bcef74 100644
--- a/examples/package.json
+++ b/examples/package.json
@@ -13,6 +13,7 @@
     "@types/node": "^24.10.0"
   },
   "dependencies": {
+    "@ai-sdk/google": "^3.0.79",
     "@ai-sdk/openai": "^3.0.65",
     "@scrawn/analytics": "link:@scrawn/analytics",
     "@scrawn/core": "link:@scrawn/core",

From 08c4923485829a40450c4cecbf441f874e3eed9d Mon Sep 17 00:00:00 2001
From: Devyash Saini <dysaini2004@gmail.com>
Date: Thu, 28 May 2026 11:04:30 +0530
Subject: [PATCH 09/11] fix(ai): fix critical operator precedence bug, remove
 double-billing, clean up types
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Fix cacheTokens operator precedence (?? vs !==) — input/outputCacheTokens
  now correctly assigned from usage, not from overrides
- Remove onFinish double-billing — per-step via onStepFinish covers usage fully
- Remove dead StripScrawnParams type
- Make BillableCallParams.userId optional — no wasted '' default
---
 packages/scrawn/src/core/ai/track.ts | 13 +++----------
 packages/scrawn/src/core/ai/types.ts | 10 ++--------
 packages/scrawn/src/core/ai/wrap.ts  | 12 +-----------
 3 files changed, 6 insertions(+), 29 deletions(-)

diff --git a/packages/scrawn/src/core/ai/track.ts b/packages/scrawn/src/core/ai/track.ts
index ec3a4b2..90cdf7c 100644
--- a/packages/scrawn/src/core/ai/track.ts
+++ b/packages/scrawn/src/core/ai/track.ts
@@ -29,17 +29,10 @@ export function buildAIPayload<TTag extends string = string>(
     outputTokens: usage.outputTokens,
     inputDebit: overrides.inputDebit ?? defaults.inputDebit,
     outputDebit: overrides.outputDebit ?? defaults.outputDebit,
-    provider:
-      (overrides.provider as string) ?? defaults.provider ?? model.provider,
-    inputCacheTokens:
-      overrides.inputCacheDebit ?? usage.inputCachedTokens !== undefined
-        ? usage.inputCachedTokens
-        : undefined,
+    provider: overrides.provider ?? defaults.provider ?? model.provider,
+    inputCacheTokens: usage.inputCachedTokens,
     inputCacheDebit: overrides.inputCacheDebit ?? defaults.inputCacheDebit,
-    outputCacheTokens:
-      overrides.outputCacheDebit ?? usage.outputCachedTokens !== undefined
-        ? usage.outputCachedTokens
-        : undefined,
+    outputCacheTokens: usage.outputCachedTokens,
     outputCacheDebit: overrides.outputCacheDebit ?? defaults.outputCacheDebit,
   };
 }
diff --git a/packages/scrawn/src/core/ai/types.ts b/packages/scrawn/src/core/ai/types.ts
index a80ef93..f07621e 100644
--- a/packages/scrawn/src/core/ai/types.ts
+++ b/packages/scrawn/src/core/ai/types.ts
@@ -22,8 +22,8 @@ export interface BillableAIOptions<TTag extends string = string> {
  * All original AI SDK params are preserved; only userId is added.
  */
 export interface BillableCallParams<TTag extends string = string> {
-  /** The user ID to bill against. */
-  userId: string;
+  /** The user ID to bill against. If omitted, billing is skipped. */
+  userId?: string;
   /** Override input token billing for this specific call. */
   inputDebit?: DebitField<TTag>;
   /** Override output token billing for this specific call. */
@@ -36,12 +36,6 @@ export interface BillableCallParams<TTag extends string = string> {
   provider?: string;
 }
 
-/**
- * Strips the BillableCallParams from a type, leaving only original AI SDK params.
- * Used internally to forward the original params to the real AI SDK.
- */
-export type StripScrawnParams<T> = Omit<T, keyof BillableCallParams>;
-
 /**
  * Language model usage as returned by Vercel AI SDK event listeners.
  * Matches the shape of OnStepFinishEvent.usage and OnFinishEvent.totalUsage.
diff --git a/packages/scrawn/src/core/ai/wrap.ts b/packages/scrawn/src/core/ai/wrap.ts
index a1e6bbb..509d08d 100644
--- a/packages/scrawn/src/core/ai/wrap.ts
+++ b/packages/scrawn/src/core/ai/wrap.ts
@@ -2,7 +2,6 @@ import type { Scrawn } from "../scrawn.js";
 import type {
   BillableAIOptions,
   BillableCallParams,
-  StripScrawnParams,
   ModelInfo,
 } from "./types.js";
 
@@ -51,7 +50,7 @@ export function createBillableAI<TTag extends string>(
         return original.apply(sdk, args);
       }
 
-      const { onStepFinish: userStep, onFinish: userFinish, ...rest } = params;
+      const { onStepFinish: userStep, ...rest } = params;
       const billingParams = { ...rest };
 
       const defaults = {
@@ -102,14 +101,6 @@ export function createBillableAI<TTag extends string>(
           userStep as ((e: unknown) => void) | undefined
         );
       }
-      if (typeof userFinish === "function") {
-        billingParams.onFinish = (event: unknown) => {
-          billingStep(
-            event as { model: ModelInfo; usage: Record<string, unknown> }
-          );
-          (userFinish as (e: unknown) => void)(event);
-        };
-      }
 
       return original.call(sdk, billingParams);
     };
@@ -122,7 +113,6 @@ function extractBillingParams<TTag extends string>(
   params: Record<string, unknown>
 ): BillableCallParams<TTag> {
   return {
-    userId: "",
     inputDebit: params.inputDebit as BillableCallParams<TTag>["inputDebit"],
     outputDebit: params.outputDebit as BillableCallParams<TTag>["outputDebit"],
     inputCacheDebit:

From 801bfd30ad333defeb150579c28d442c0e7afc2d Mon Sep 17 00:00:00 2001
From: Devyash Saini <dysaini2004@gmail.com>
Date: Thu, 28 May 2026 11:13:14 +0530
Subject: [PATCH 10/11] =?UTF-8?q?refactor:=20fix=20all=206=20fallow=20find?=
 =?UTF-8?q?ings=20=E2=80=94=20remove=20dead=20type,=20extract=20duplicatio?=
 =?UTF-8?q?n?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

1. Remove unused TagName type from pricing/types.ts
2. Add formatValidationError helper — removes duplicate Zod issue formatting (3 sites)
3. Add performAIStreamCall helper — removes duplicate gRPC stream call (2 sites)
4. Extract initClient/buildCallOptions into grpc/utils.ts — shared by both builders
5. Replace duplicate client init in requestBuilder + streamRequestBuilder with shared utils
6. Extract isValidPriceExpr helper — deduplicates null/type/kind check in both schemas
---
 .../scrawn/src/core/grpc/requestBuilder.ts    |  13 +--
 .../src/core/grpc/streamRequestBuilder.ts     |  13 +--
 packages/scrawn/src/core/grpc/utils.ts        |  34 ++++++
 packages/scrawn/src/core/pricing/types.ts     |   9 --
 packages/scrawn/src/core/scrawn.ts            | 100 +++++++++---------
 packages/scrawn/src/core/types/event.ts       |  48 ++++-----
 6 files changed, 111 insertions(+), 106 deletions(-)
 create mode 100644 packages/scrawn/src/core/grpc/utils.ts

diff --git a/packages/scrawn/src/core/grpc/requestBuilder.ts b/packages/scrawn/src/core/grpc/requestBuilder.ts
index e0f4686..88925ec 100644
--- a/packages/scrawn/src/core/grpc/requestBuilder.ts
+++ b/packages/scrawn/src/core/grpc/requestBuilder.ts
@@ -1,6 +1,7 @@
 import * as grpc from "@grpc/grpc-js";
 import type { GrpcCallOptions } from "./types.js";
 import type { GrpcCallContext } from "./callContext.js";
+import { initClient, buildCallOptions, getRequestMetadata } from "./utils.js";
 
 export class RequestBuilder<
   C extends { new (...args: any[]): any; serviceName: string }
@@ -47,10 +48,7 @@ export class RequestBuilder<
     this.ctx.logCallStart();
 
     try {
-      const client = new this.ctx.ClientConstructor(
-        this.ctx.target,
-        this.ctx.credentials
-      ) as grpc.Client & Record<string, unknown>;
+      const client = initClient(this.ctx);
       const method = client[this.ctx.methodName] as (
         request: unknown,
         metadata: grpc.Metadata,
@@ -58,16 +56,13 @@ export class RequestBuilder<
         callback: (error: grpc.ServiceError | null, response: TResponse) => void
       ) => void;
 
-      const callOptions: grpc.CallOptions = {};
-      if (this.options.deadline !== undefined) {
-        callOptions.deadline = this.options.deadline;
-      }
+      const callOptions = buildCallOptions(this.options);
 
       const response = await new Promise<TResponse>((resolve, reject) => {
         method.call(
           client,
           this.payload,
-          this.options.metadata ?? this.ctx.getMetadata(),
+          getRequestMetadata(this.options, this.ctx),
           callOptions,
           (error, response) => {
             if (error) {
diff --git a/packages/scrawn/src/core/grpc/streamRequestBuilder.ts b/packages/scrawn/src/core/grpc/streamRequestBuilder.ts
index bac9d2c..0c37129 100644
--- a/packages/scrawn/src/core/grpc/streamRequestBuilder.ts
+++ b/packages/scrawn/src/core/grpc/streamRequestBuilder.ts
@@ -1,6 +1,7 @@
 import * as grpc from "@grpc/grpc-js";
 import type { GrpcCallOptions } from "./types.js";
 import type { GrpcCallContext } from "./callContext.js";
+import { initClient, buildCallOptions, getRequestMetadata } from "./utils.js";
 
 export class StreamRequestBuilder<
   C extends { new (...args: any[]): any; serviceName: string }
@@ -37,25 +38,19 @@ export class StreamRequestBuilder<
     this.ctx.logCallStart();
 
     try {
-      const client = new this.ctx.ClientConstructor(
-        this.ctx.target,
-        this.ctx.credentials
-      ) as grpc.Client & Record<string, unknown>;
+      const client = initClient(this.ctx);
       const method = client[this.ctx.methodName] as (
         metadata: grpc.Metadata,
         options: grpc.CallOptions,
         callback: (error: grpc.ServiceError | null, response: TResponse) => void
       ) => grpc.ClientWritableStream<unknown>;
 
-      const callOptions: grpc.CallOptions = {};
-      if (this.options.deadline !== undefined) {
-        callOptions.deadline = this.options.deadline;
-      }
+      const callOptions = buildCallOptions(this.options);
 
       const response = await new Promise<TResponse>((resolve, reject) => {
         const stream = method.call(
           client,
-          this.options.metadata ?? this.ctx.getMetadata(),
+          getRequestMetadata(this.options, this.ctx),
           callOptions,
           (error, result) => {
             if (error) {
diff --git a/packages/scrawn/src/core/grpc/utils.ts b/packages/scrawn/src/core/grpc/utils.ts
new file mode 100644
index 0000000..01fe71b
--- /dev/null
+++ b/packages/scrawn/src/core/grpc/utils.ts
@@ -0,0 +1,34 @@
+import * as grpc from "@grpc/grpc-js";
+import type { GrpcCallOptions } from "./types.js";
+import type { GrpcCallContext } from "./callContext.js";
+
+/**
+ * Shared client initializer used by both RequestBuilder and StreamRequestBuilder.
+ */
+export function initClient<
+  C extends { new (...args: any[]): any; serviceName: string }
+>(ctx: GrpcCallContext<C>): grpc.Client & Record<string, unknown> {
+  return new ctx.ClientConstructor(ctx.target, ctx.credentials) as grpc.Client &
+    Record<string, unknown>;
+}
+
+/**
+ * Shared call-options builder from optional GrpcCallOptions.
+ */
+export function buildCallOptions(options: GrpcCallOptions): grpc.CallOptions {
+  const callOptions: grpc.CallOptions = {};
+  if (options.deadline !== undefined) {
+    callOptions.deadline = options.deadline;
+  }
+  return callOptions;
+}
+
+/**
+ * Shared metadata getter — prefers per-request metadata, falls back to context metadata.
+ */
+export function getRequestMetadata(
+  options: GrpcCallOptions,
+  ctx: { getMetadata(): grpc.Metadata }
+): grpc.Metadata {
+  return options.metadata ?? ctx.getMetadata();
+}
diff --git a/packages/scrawn/src/core/pricing/types.ts b/packages/scrawn/src/core/pricing/types.ts
index cb4bf96..4f7427c 100644
--- a/packages/scrawn/src/core/pricing/types.ts
+++ b/packages/scrawn/src/core/pricing/types.ts
@@ -19,15 +19,6 @@
  */
 export type OpType = "ADD" | "SUB" | "MUL" | "DIV";
 
-/**
- * Intellisense hint type for tag names.
- * Tag names must be ALL CAPS with underscores only (e.g., PREMIUM_CALL, FEE, INPUT_RATE).
- * No lowercase, digits, or hyphens allowed.
- *
- * This is a branded type that provides IDE hints while remaining compatible with `string`.
- */
-export type TagName = Uppercase<string> & { readonly __brand?: "TagName" };
-
 /**
  * A literal amount in cents (must be an integer).
  */
diff --git a/packages/scrawn/src/core/scrawn.ts b/packages/scrawn/src/core/scrawn.ts
index e008fc2..7ed87be 100644
--- a/packages/scrawn/src/core/scrawn.ts
+++ b/packages/scrawn/src/core/scrawn.ts
@@ -154,6 +154,24 @@ export class Scrawn<
     return error;
   }
 
+  /** Shared: formats Zod issues into a ScrawnValidationError and notifies the callback. */
+  private formatValidationError(
+    message: string,
+    issues: import("zod").ZodIssue[],
+    onError?: EventConsumerErrorCallback
+  ): ScrawnValidationError {
+    const error = new ScrawnValidationError(message, {
+      details: {
+        errors: issues.map((e) => ({
+          field: e.path.join("."),
+          message: e.message,
+        })),
+      },
+    });
+    this.notifyValidationError(error, onError);
+    return error;
+  }
+
   /**
    * Creates a new Scrawn SDK instance.
    *
@@ -398,15 +416,11 @@ export class Scrawn<
         .map((e) => `${e.path.join(".")}: ${e.message}`)
         .join(", ");
       log.error(`Invalid payload for basicUsageEventConsumer: ${errors}`);
-      const error = new ScrawnValidationError("Payload validation failed", {
-        details: {
-          errors: validationResult.error.issues.map((e) => ({
-            field: e.path.join("."),
-            message: e.message,
-          })),
-        },
-      });
-      this.notifyValidationError(error, options?.onError);
+      this.formatValidationError(
+        "Payload validation failed",
+        validationResult.error.issues,
+        options?.onError
+      );
       return;
     }
 
@@ -565,15 +579,11 @@ export class Scrawn<
           log.error(
             `Invalid payload extracted in middlewareEventConsumer: ${errors}`
           );
-          const error = new ScrawnValidationError("Payload validation failed", {
-            details: {
-              errors: validationResult.error.issues.map((e) => ({
-                field: e.path.join("."),
-                message: e.message,
-              })),
-            },
-          });
-          this.notifyValidationError(error, config.onError);
+          this.formatValidationError(
+            "Payload validation failed",
+            validationResult.error.issues,
+            config.onError
+          );
           return next();
         }
 
@@ -941,27 +951,12 @@ export class Scrawn<
       const responsePromise = (async (): Promise<
         StreamEventResponse | undefined
       > => {
-        try {
-          log.info("Starting AI token usage stream (return mode)");
-
-          const response = await this.grpcClient
-            .newStreamCall(EventServiceClient, "streamEvents")
-            .addMetadata("authorization", `Bearer ${creds.apiKey}`)
-            .stream<StreamEventResponse>(transformedStream);
-
-          log.info(
-            `AI token stream completed: ${response.eventsProcessed} events processed`
-          );
-          return response;
-        } catch (error) {
-          log.error(
-            `Failed to stream AI token usage: ${
-              error instanceof Error ? error.message : "Unknown error"
-            }`
-          );
-          this.notifyEventConsumerError(error, onError);
-          return undefined;
-        }
+        const result = await this.performAIStreamCall(
+          creds.apiKey,
+          transformedStream,
+          onError
+        );
+        return result;
       })();
 
       return { response: responsePromise, stream: userStream };
@@ -970,12 +965,24 @@ export class Scrawn<
     // Default: fire-and-forget mode
     const transformedStream = this.transformAITokenStream(stream, onError);
 
+    return this.performAIStreamCall(creds.apiKey, transformedStream, onError);
+  }
+
+  /**
+   * Shared: performs a gRPC streaming call for AI token events.
+   * Used by both return-mode and fire-and-forget branches of aiTokenStreamConsumer.
+   */
+  private async performAIStreamCall(
+    apiKey: string,
+    transformedStream: AsyncIterable<unknown>,
+    onError?: EventConsumerErrorCallback
+  ): Promise<StreamEventResponse | undefined> {
     try {
       log.info("Starting AI token usage stream");
 
       const response = await this.grpcClient
         .newStreamCall(EventServiceClient, "streamEvents")
-        .addMetadata("authorization", `Bearer ${creds.apiKey}`)
+        .addMetadata("authorization", `Bearer ${apiKey}`)
         .stream<StreamEventResponse>(transformedStream);
 
       log.info(
@@ -1051,18 +1058,11 @@ export class Scrawn<
           .map((e) => `${e.path.join(".")}: ${e.message}`)
           .join(", ");
         log.error(`Invalid AI token usage payload, skipping: ${errors}`);
-        const error = new ScrawnValidationError(
+        this.formatValidationError(
           "AI token usage payload validation failed",
-          {
-            details: {
-              errors: validationResult.error.issues.map((e) => ({
-                field: e.path.join("."),
-                message: e.message,
-              })),
-            },
-          }
+          validationResult.error.issues,
+          onError
         );
-        this.notifyValidationError(error, onError);
         continue;
       }
 
diff --git a/packages/scrawn/src/core/types/event.ts b/packages/scrawn/src/core/types/event.ts
index 5c6629a..50e93c8 100644
--- a/packages/scrawn/src/core/types/event.ts
+++ b/packages/scrawn/src/core/types/event.ts
@@ -16,23 +16,28 @@ const ALL_EXPR_KINDS = [
   "exprRef",
 ];
 
+/**
+ * Custom zod schema for PriceExpr validation (allows token placeholders).
+ * Used in AI token usage payloads where inputTokens()/outputTokens() are valid.
+ */
+/**
+ * Shared validator: checks the value is a valid PriceExpr structure (kind exists, expr validates).
+ * Does NOT check for token placeholders — that's the caller's responsibility.
+ */
+function isValidPriceExpr(val: unknown): val is PriceExpr<string> {
+  if (val === null || val === undefined || typeof val !== "object") {
+    return false;
+  }
+  const expr = val as PriceExpr<string>;
+  return ALL_EXPR_KINDS.includes(expr.kind) && isValidExpr(expr);
+}
+
 /**
  * Custom zod schema for PriceExpr validation (allows token placeholders).
  * Used in AI token usage payloads where inputTokens()/outputTokens() are valid.
  */
 const PriceExprSchema = z.custom<PriceExpr<string>>(
-  (val): val is PriceExpr<string> => {
-    if (val === null || val === undefined || typeof val !== "object") {
-      return false;
-    }
-    const expr = val as PriceExpr<string>;
-    // Check that it has a valid kind (including token placeholders)
-    if (!ALL_EXPR_KINDS.includes(expr.kind)) {
-      return false;
-    }
-    // Use the validation function
-    return isValidExpr(expr);
-  },
+  (val): val is PriceExpr<string> => isValidPriceExpr(val),
   {
     message:
       "Must be a valid pricing expression (use tag(), add(), sub(), mul(), div(), amount(), inputTokens(), or outputTokens())",
@@ -46,23 +51,8 @@ const PriceExprSchema = z.custom<PriceExpr<string>>(
  */
 const PriceExprNoTokensSchema = z.custom<PriceExpr<string>>(
   (val): val is PriceExpr<string> => {
-    if (val === null || val === undefined || typeof val !== "object") {
-      return false;
-    }
-    const expr = val as PriceExpr<string>;
-    // Check that it has a valid kind
-    if (!ALL_EXPR_KINDS.includes(expr.kind)) {
-      return false;
-    }
-    // Use the validation function
-    if (!isValidExpr(expr)) {
-      return false;
-    }
-    // Reject token placeholders in SDK call context
-    if (containsTokenExpr(expr)) {
-      return false;
-    }
-    return true;
+    if (!isValidPriceExpr(val)) return false;
+    return !containsTokenExpr(val as PriceExpr<string>);
   },
   {
     message:

From c197061b55d6bb5f1ee776ebd5d21e8b05b6cff0 Mon Sep 17 00:00:00 2001
From: Devyash Saini <dysaini2004@gmail.com>
Date: Thu, 28 May 2026 11:19:55 +0530
Subject: [PATCH 11/11] chore: update proto submodule ref

---
 packages/scrawn/proto | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/packages/scrawn/proto b/packages/scrawn/proto
index 8e70696..d759346 160000
--- a/packages/scrawn/proto
+++ b/packages/scrawn/proto
@@ -1 +1 @@
-Subproject commit 8e706964fedccfe107152d2ec0063acae7f2b35d
+Subproject commit d759346676331bc1bb15f6d40f6ef3dd1aeb8706