diff --git a/packages/scrawn/proto/event/v1/event.proto b/packages/scrawn/proto/event/v1/event.proto index f2a7296..9f82812 100644 --- a/packages/scrawn/proto/event/v1/event.proto +++ b/packages/scrawn/proto/event/v1/event.proto @@ -36,6 +36,7 @@ message SDKCall { oneof debit { float amount = 2; string tag = 3; + string expr = 4; // Pricing expression (e.g., "add(mul(tag('PREMIUM'),3),250)") } } @@ -60,11 +61,13 @@ message AITokenUsage { oneof inputDebit { float inputAmount = 4; string inputTag = 5; + string inputExpr = 8; // Pricing expression for input tokens } oneof outputDebit { float outputAmount = 6; string outputTag = 7; + string outputExpr = 9; // Pricing expression for output tokens } } diff --git a/packages/scrawn/src/core/grpc/requestBuilder.ts b/packages/scrawn/src/core/grpc/requestBuilder.ts index 2e1be7a..7760c04 100644 --- a/packages/scrawn/src/core/grpc/requestBuilder.ts +++ b/packages/scrawn/src/core/grpc/requestBuilder.ts @@ -27,6 +27,7 @@ import type { MethodInput, MethodOutput, RequestState, + UnaryMethodFn, } from './types.js'; import type { GrpcCallContext } from './callContext.js'; @@ -152,14 +153,20 @@ export class RequestBuilder< this.ctx.logCallStart(); this.ctx.log.debug(`Payload: ${JSON.stringify(this.payload)}`); - // The actual gRPC call - fully type-safe! - const response = await (this.ctx.client[this.ctx.methodName] as any)( - this.payload, - { headers: this.ctx.getHeaders() } - ); + // The actual gRPC call. + // Type assertion is required because TypeScript cannot track the relationship + // between a dynamic key access and the resulting function signature in mapped types. + // This is safe because: + // 1. methodName is constrained to ServiceMethodNames + // 2. MethodInput/MethodOutput are derived from the same service definition + const method = this.ctx.client[this.ctx.methodName] as UnaryMethodFn< + MethodInput, + MethodOutput + >; + const response = await method(this.payload, { headers: this.ctx.getHeaders() }); this.ctx.logCallSuccess(); - return response as MethodOutput; + return response; } catch (error) { this.ctx.logCallError(error); throw error; diff --git a/packages/scrawn/src/core/grpc/streamRequestBuilder.ts b/packages/scrawn/src/core/grpc/streamRequestBuilder.ts index 6988de0..61473c8 100644 --- a/packages/scrawn/src/core/grpc/streamRequestBuilder.ts +++ b/packages/scrawn/src/core/grpc/streamRequestBuilder.ts @@ -21,6 +21,7 @@ import type { ServiceMethodNames, MethodInput, MethodOutput, + ClientStreamingMethodFn, } from './types.js'; import type { GrpcCallContext } from './callContext.js'; @@ -101,14 +102,23 @@ export class StreamRequestBuilder< try { this.ctx.logCallStart(); - // The actual client-streaming gRPC call - const response = await (this.ctx.client[this.ctx.methodName] as any)( - iterable, + // The actual client-streaming gRPC call. + // Type assertion is required because TypeScript cannot track the relationship + // between a dynamic key access and the resulting function signature in mapped types. + // This is safe because: + // 1. methodName is constrained to ServiceMethodNames + // 2. MethodInput/MethodOutput are derived from the same service definition + const method = this.ctx.client[this.ctx.methodName] as ClientStreamingMethodFn< + MethodInput, + MethodOutput + >; + const response = await method( + iterable as AsyncIterable>>, { headers: this.ctx.getHeaders() } ); this.ctx.logCallSuccess(); - return response as MethodOutput; + return response; } catch (error) { this.ctx.logCallError(error); throw error; diff --git a/packages/scrawn/src/core/grpc/types.ts b/packages/scrawn/src/core/grpc/types.ts index 8dad68d..9f11bfb 100644 --- a/packages/scrawn/src/core/grpc/types.ts +++ b/packages/scrawn/src/core/grpc/types.ts @@ -115,3 +115,61 @@ export type ServerStreamingMethodNames = MethodsOfKind = MethodsOfKind; + +/** + * Options passed to gRPC method calls. + * Matches CallOptions from @connectrpc/connect. + */ +export interface GrpcCallOptions { + headers?: Headers; + signal?: AbortSignal; + timeoutMs?: number; + onHeader?: (headers: globalThis.Headers) => void; + onTrailer?: (trailers: globalThis.Headers) => void; +} + +/** + * Function signature for a unary gRPC method. + * Takes a partial message and options, returns a promise of the response. + * + * @template I - Input message type + * @template O - Output message type + */ +export type UnaryMethodFn = ( + request: Partial, + options?: GrpcCallOptions +) => Promise; + +/** + * Function signature for a client-streaming gRPC method. + * Takes an async iterable of partial messages and options, returns a promise of the response. + * + * @template I - Input message type + * @template O - Output message type + */ +export type ClientStreamingMethodFn = ( + request: AsyncIterable>, + options?: GrpcCallOptions +) => Promise; + +/** + * Get the method function type from a Client for a specific method. + * This properly extracts the callable type for dynamic method access. + * + * Note: Due to TypeScript limitations with mapped types and dynamic keys, + * this type is used with type assertions when accessing client methods + * via computed property names. The assertion is safe because: + * 1. The method name M is constrained to ServiceMethodNames + * 2. The input/output types are derived from the same service definition + * + * @template S - The gRPC service type + * @template M - The method name on the service + */ +export type ClientMethod< + S extends ServiceType, + M extends ServiceMethodNames +> = S['methods'][M] extends { kind: 'unary' } + ? UnaryMethodFn, MethodOutput> + : S['methods'][M] extends { kind: 'client_streaming' } + ? ClientStreamingMethodFn, MethodOutput> + : never; diff --git a/packages/scrawn/src/core/pricing/builders.ts b/packages/scrawn/src/core/pricing/builders.ts new file mode 100644 index 0000000..1a7e321 --- /dev/null +++ b/packages/scrawn/src/core/pricing/builders.ts @@ -0,0 +1,167 @@ +/** + * Pricing DSL Builder Functions + * + * This module provides the fluent builder functions for constructing + * type-safe pricing expressions. These functions perform light validation + * and build the AST that gets serialized to a string for the backend. + * + * @example + * ```typescript + * import { add, mul, tag } from '@scrawn/core'; + * + * // (PREMIUM_CALL * 3) + EXTRA_FEE + 250 cents + * const expr = add(mul(tag('PREMIUM_CALL'), 3), tag('EXTRA_FEE'), 250); + * ``` + */ + +import type { AmountExpr, TagExpr, OpExpr, PriceExpr, ExprInput } from './types.js'; +import { validateExpr } from './validate.js'; + +/** + * Convert an ExprInput (PriceExpr or number) to a PriceExpr. + * Numbers are wrapped as AmountExpr (cents). + */ +function toExpr(input: ExprInput): PriceExpr { + if (typeof input === 'number') { + return { kind: 'amount', value: input } as const; + } + return input; +} + +/** + * Create a tag reference expression. + * Tags are resolved to their cent values by the backend. + * + * @param name - The name of the price tag (must be non-empty) + * @returns A TagExpr referencing the named tag + * @throws Error if name is empty or whitespace-only + * + * @example + * ```typescript + * const premiumTag = tag('PREMIUM_CALL'); + * ``` + */ +export function tag(name: string): TagExpr { + const expr: TagExpr = { kind: 'tag', name } as const; + validateExpr(expr); // Will throw if invalid + return expr; +} + +/** + * Create an addition expression. + * Adds all arguments together: arg1 + arg2 + arg3 + ... + * + * @param args - Two or more expressions or numbers (cents) to add + * @returns An OpExpr representing the sum + * @throws Error if fewer than 2 arguments provided + * + * @example + * ```typescript + * // 100 + 200 + tag('BONUS') + * const sum = add(100, 200, tag('BONUS')); + * ``` + */ +export function add(...args: ExprInput[]): OpExpr { + const expr: OpExpr = { + kind: 'op', + op: 'ADD', + args: args.map(toExpr), + } as const; + validateExpr(expr); + return expr; +} + +/** + * Create a subtraction expression. + * Subtracts subsequent arguments from the first: arg1 - arg2 - arg3 - ... + * + * @param args - Two or more expressions or numbers (cents) to subtract + * @returns An OpExpr representing the difference + * @throws Error if fewer than 2 arguments provided + * + * @example + * ```typescript + * // tag('TOTAL') - 50 + * const diff = sub(tag('TOTAL'), 50); + * ``` + */ +export function sub(...args: ExprInput[]): OpExpr { + const expr: OpExpr = { + kind: 'op', + op: 'SUB', + args: args.map(toExpr), + } as const; + validateExpr(expr); + return expr; +} + +/** + * Create a multiplication expression. + * Multiplies all arguments together: arg1 * arg2 * arg3 * ... + * + * @param args - Two or more expressions or numbers (cents) to multiply + * @returns An OpExpr representing the product + * @throws Error if fewer than 2 arguments provided + * + * @example + * ```typescript + * // tag('PER_TOKEN') * 100 + * const product = mul(tag('PER_TOKEN'), 100); + * ``` + */ +export function mul(...args: ExprInput[]): OpExpr { + const expr: OpExpr = { + kind: 'op', + op: 'MUL', + args: args.map(toExpr), + } as const; + validateExpr(expr); + return expr; +} + +/** + * Create a division expression. + * Divides the first argument by subsequent arguments: arg1 / arg2 / arg3 / ... + * + * Note: The backend performs integer division. Results are truncated, not rounded. + * + * @param args - Two or more expressions or numbers (cents) to divide + * @returns An OpExpr representing the quotient + * @throws Error if fewer than 2 arguments provided + * @throws Error if any literal divisor is 0 (detected at build time) + * + * @example + * ```typescript + * // tag('TOTAL') / 2 + * const half = div(tag('TOTAL'), 2); + * ``` + */ +export function div(...args: ExprInput[]): OpExpr { + const expr: OpExpr = { + kind: 'op', + op: 'DIV', + args: args.map(toExpr), + } as const; + validateExpr(expr); + return expr; +} + +/** + * Create an amount expression from a number of cents. + * This is useful when you need to explicitly create an AmountExpr + * rather than relying on auto-conversion. + * + * @param cents - The amount in cents (must be an integer) + * @returns An AmountExpr representing the amount + * @throws Error if cents is not a finite integer + * + * @example + * ```typescript + * const fee = amount(250); // 250 cents = $2.50 + * ``` + */ +export function amount(cents: number): AmountExpr { + const expr: AmountExpr = { kind: 'amount', value: cents } as const; + validateExpr(expr); + return expr; +} diff --git a/packages/scrawn/src/core/pricing/index.ts b/packages/scrawn/src/core/pricing/index.ts new file mode 100644 index 0000000..183794a --- /dev/null +++ b/packages/scrawn/src/core/pricing/index.ts @@ -0,0 +1,42 @@ +/** + * Pricing DSL Module + * + * Provides a type-safe DSL for building pricing expressions that combine + * literal amounts, named price tags, and arithmetic operations. + * + * @example + * ```typescript + * import { add, mul, tag, serializeExpr } from '@scrawn/core'; + * + * // Build expression: (PREMIUM_CALL * 3) + EXTRA_FEE + 250 cents + * const expr = add(mul(tag('PREMIUM_CALL'), 3), tag('EXTRA_FEE'), 250); + * + * // Serialize for backend: "add(mul(tag('PREMIUM_CALL'),3),tag('EXTRA_FEE'),250)" + * const exprString = serializeExpr(expr); + * + * // Use in event payload + * await scrawn.sdkCallEventConsumer({ + * userId: 'u123', + * debitExpr: expr + * }); + * ``` + */ + +// Export types +export type { + OpType, + AmountExpr, + TagExpr, + OpExpr, + PriceExpr, + ExprInput, +} from './types.js'; + +// Export builder functions +export { tag, add, sub, mul, div, amount } from './builders.js'; + +// Export serialization +export { serializeExpr, prettyPrintExpr } from './serialize.js'; + +// Export validation +export { validateExpr, isValidExpr, PricingExpressionError } from './validate.js'; diff --git a/packages/scrawn/src/core/pricing/serialize.ts b/packages/scrawn/src/core/pricing/serialize.ts new file mode 100644 index 0000000..2649396 --- /dev/null +++ b/packages/scrawn/src/core/pricing/serialize.ts @@ -0,0 +1,123 @@ +/** + * Pricing DSL Serialization + * + * This module converts typed pricing expression ASTs into string format + * that the backend can parse and evaluate. + * + * Output format examples: + * - Amount: 250 + * - Tag: tag('PREMIUM_CALL') + * - Addition: add(100,tag('FEE'),250) + * - Complex: add(mul(tag('PREMIUM_CALL'),3),tag('EXTRA_FEE'),250) + * + * The format is designed to be: + * - Unambiguous (parseable by the backend) + * - Human-readable (for debugging) + * - Compact (minimal whitespace) + */ + +import type { PriceExpr, AmountExpr, TagExpr, OpExpr } from './types.js'; + +/** + * Serialize a pricing expression to a string. + * + * @param expr - The expression to serialize + * @returns A string representation that the backend can parse + * + * @example + * ```typescript + * const expr = add(mul(tag('PREMIUM'), 3), 100); + * const str = serializeExpr(expr); + * // "add(mul(tag('PREMIUM'),3),100)" + * ``` + */ +export function serializeExpr(expr: PriceExpr): string { + switch (expr.kind) { + case 'amount': + return serializeAmount(expr); + case 'tag': + return serializeTag(expr); + case 'op': + return serializeOp(expr); + } +} + +/** + * Serialize an amount expression. + * Just the integer value. + */ +function serializeAmount(expr: AmountExpr): string { + return expr.value.toString(); +} + +/** + * Serialize a tag expression. + * Format: tag('TAG_NAME') + * Uses single quotes for the tag name. + */ +function serializeTag(expr: TagExpr): string { + // Escape single quotes in tag name (though validation should prevent this) + const escapedName = expr.name.replace(/'/g, "\\'"); + return `tag('${escapedName}')`; +} + +/** + * Serialize an operation expression. + * Format: op(arg1,arg2,arg3,...) + */ +function serializeOp(expr: OpExpr): string { + const opName = expr.op.toLowerCase(); + const args = expr.args.map(serializeExpr).join(','); + return `${opName}(${args})`; +} + +/** + * Pretty-print a pricing expression with indentation. + * Useful for debugging and logging. + * + * @param expr - The expression to format + * @param indent - Number of spaces for indentation (default 2) + * @returns A formatted, multi-line string representation + * + * @example + * ```typescript + * const expr = add(mul(tag('PREMIUM'), 3), 100); + * console.log(prettyPrintExpr(expr)); + * // add( + * // mul( + * // tag('PREMIUM'), + * // 3 + * // ), + * // 100 + * // ) + * ``` + */ +export function prettyPrintExpr(expr: PriceExpr, indent: number = 2): string { + return prettyPrintInternal(expr, 0, indent); +} + +function prettyPrintInternal(expr: PriceExpr, level: number, indent: number): string { + const pad = ' '.repeat(level * indent); + + switch (expr.kind) { + case 'amount': + return expr.value.toString(); + case 'tag': + return `tag('${expr.name}')`; + case 'op': { + const opName = expr.op.toLowerCase(); + if (expr.args.length === 0) { + return `${opName}()`; + } + + const args = expr.args + .map(arg => { + const inner = prettyPrintInternal(arg, level + 1, indent); + return ' '.repeat((level + 1) * indent) + inner; + }) + .join(',\n'); + + return `${opName}(\n${args}\n${pad})`; + } + } +} diff --git a/packages/scrawn/src/core/pricing/types.ts b/packages/scrawn/src/core/pricing/types.ts new file mode 100644 index 0000000..3323eaa --- /dev/null +++ b/packages/scrawn/src/core/pricing/types.ts @@ -0,0 +1,56 @@ +/** + * Pricing DSL Types + * + * This module defines the type-safe AST for pricing expressions. + * The SDK builds typed expressions using these types, then serializes + * them to strings for the backend to parse and evaluate. + * + * @example + * ```typescript + * import { add, mul, tag } from '@scrawn/core'; + * + * // Build a pricing expression: (PREMIUM_CALL * 3) + EXTRA_FEE + 250 cents + * const expr = add(mul(tag('PREMIUM_CALL'), 3), tag('EXTRA_FEE'), 250); + * ``` + */ + +/** + * Supported arithmetic operations for pricing expressions. + */ +export type OpType = 'ADD' | 'SUB' | 'MUL' | 'DIV'; + +/** + * A literal amount in cents (must be an integer). + */ +export interface AmountExpr { + readonly kind: 'amount'; + readonly value: number; +} + +/** + * A reference to a named price tag (resolved by the backend). + */ +export interface TagExpr { + readonly kind: 'tag'; + readonly name: string; +} + +/** + * An arithmetic operation combining multiple expressions. + */ +export interface OpExpr { + readonly kind: 'op'; + readonly op: OpType; + readonly args: readonly PriceExpr[]; +} + +/** + * A pricing expression - can be a literal amount, a tag reference, or an operation. + */ +export type PriceExpr = AmountExpr | TagExpr | OpExpr; + +/** + * Input type for DSL builder functions. + * Accepts either a PriceExpr or a raw number (interpreted as cents). + */ +export type ExprInput = PriceExpr | number; diff --git a/packages/scrawn/src/core/pricing/validate.ts b/packages/scrawn/src/core/pricing/validate.ts new file mode 100644 index 0000000..ee335b2 --- /dev/null +++ b/packages/scrawn/src/core/pricing/validate.ts @@ -0,0 +1,151 @@ +/** + * Pricing DSL Validation + * + * This module provides light SDK-side validation for pricing expressions. + * The backend performs full validation; the SDK only catches obvious errors + * early to provide better developer experience. + * + * SDK validates: + * - Division by literal zero + * - Non-integer cents (amounts must be integers) + * - Non-finite numbers (NaN, Infinity) + * - Empty operation arguments (ops need at least 2 args) + * - Empty/whitespace tag names + * + * SDK does NOT validate: + * - Tag existence (backend resolves tags) + * - Division by zero when divisor is a tag (backend handles) + * - Overflow (backend handles) + * - Negative results (backend handles) + */ + +import type { PriceExpr, OpExpr } from './types.js'; + +/** + * Error thrown when a pricing expression fails validation. + */ +export class PricingExpressionError extends Error { + constructor(message: string) { + super(message); + this.name = 'PricingExpressionError'; + } +} + +/** + * Validate a pricing expression. + * Throws PricingExpressionError if validation fails. + * + * @param expr - The expression to validate + * @throws PricingExpressionError if validation fails + */ +export function validateExpr(expr: PriceExpr): void { + switch (expr.kind) { + case 'amount': + validateAmount(expr.value); + break; + case 'tag': + validateTagName(expr.name); + break; + case 'op': + validateOp(expr); + break; + } +} + +/** + * Validate an amount value. + * Must be a finite integer. + */ +function validateAmount(value: number): void { + if (!Number.isFinite(value)) { + throw new PricingExpressionError( + `Amount must be a finite number, got: ${value}` + ); + } + if (!Number.isInteger(value)) { + throw new PricingExpressionError( + `Amount must be an integer (cents), got: ${value}. ` + + `Hint: Use cents instead of dollars (e.g., 250 instead of 2.50)` + ); + } +} + +/** + * Validate a tag name. + * Must be a non-empty string with no leading/trailing whitespace. + */ +function validateTagName(name: string): void { + if (typeof name !== 'string') { + throw new PricingExpressionError( + `Tag name must be a string, got: ${typeof name}` + ); + } + if (name.length === 0) { + throw new PricingExpressionError('Tag name cannot be empty'); + } + if (name.trim() !== name) { + throw new PricingExpressionError( + `Tag name cannot have leading or trailing whitespace: "${name}"` + ); + } + if (name.trim().length === 0) { + throw new PricingExpressionError('Tag name cannot be only whitespace'); + } + // Validate tag name format: alphanumeric, underscores, hyphens + if (!/^[A-Za-z_][A-Za-z0-9_-]*$/.test(name)) { + throw new PricingExpressionError( + `Tag name must start with a letter or underscore and contain only ` + + `alphanumeric characters, underscores, or hyphens: "${name}"` + ); + } +} + +/** + * Validate an operation expression. + * Must have at least 2 arguments. + * For division, checks for literal zero divisors. + */ +function validateOp(expr: OpExpr): void { + const { op, args } = expr; + + // Must have at least 2 arguments + if (args.length < 2) { + throw new PricingExpressionError( + `Operation ${op.toLowerCase()} requires at least 2 arguments, got: ${args.length}` + ); + } + + // Recursively validate all arguments + for (const arg of args) { + validateExpr(arg); + } + + // Check for division by literal zero + if (op === 'DIV') { + // Check all divisors (all args after the first) + for (let i = 1; i < args.length; i++) { + const arg = args[i]; + if (arg.kind === 'amount' && arg.value === 0) { + throw new PricingExpressionError( + `Division by zero: divisor at position ${i + 1} is 0` + ); + } + } + } +} + +/** + * Check if an expression is valid without throwing. + * Returns true if valid, false otherwise. + * + * @param expr - The expression to check + * @returns true if the expression is valid + */ +export function isValidExpr(expr: PriceExpr): boolean { + try { + validateExpr(expr); + return true; + } catch { + return false; + } +} diff --git a/packages/scrawn/src/core/scrawn.ts b/packages/scrawn/src/core/scrawn.ts index e5b4999..95e1b31 100644 --- a/packages/scrawn/src/core/scrawn.ts +++ b/packages/scrawn/src/core/scrawn.ts @@ -23,6 +23,7 @@ import { ScrawnValidationError, convertGrpcError } from './errors/index.js'; +import { serializeExpr } from './pricing/index.js'; const log = new ScrawnLogger('Scrawn'); @@ -162,34 +163,43 @@ export class Scrawn { return creds as AuthRegistry[K]; } - /** - * Track an SDK call event. - * - * Records SDK usage to the Scrawn backend for billing tracking. - * The event is authenticated using the API key provided during SDK initialization. - * - * @param payload - The SDK call data to track - * @param payload.userId - Unique identifier of the user making the call - * @param payload.debitAmount - (Optional) Direct amount to debit from the user's account - * @param payload.debitTag - (Optional) Named price tag for backend-managed pricing - * @returns A promise that resolves when the event is tracked - * @throws Error if payload validation fails or both/neither debit fields are provided - * - * @example - * ```typescript - * // Using direct amount - * await scrawn.sdkCallEventConsumer({ - * userId: 'user_abc123', - * debitAmount: 10 - * }); - * - * // Using price tag - * await scrawn.sdkCallEventConsumer({ - * userId: 'user_abc123', - * debitTag: 'PREMIUM_FEATURE' - * }); - * ``` - */ + /** + * Track an SDK call event. + * + * Records SDK usage to the Scrawn backend for billing tracking. + * The event is authenticated using the API key provided during SDK initialization. + * + * @param payload - The SDK call data to track + * @param payload.userId - Unique identifier of the user making the call + * @param payload.debitAmount - (Optional) Direct amount in cents to debit from the user's account + * @param payload.debitTag - (Optional) Named price tag for backend-managed pricing + * @param payload.debitExpr - (Optional) Pricing expression for complex calculations + * @returns A promise that resolves when the event is tracked + * @throws Error if payload validation fails or if not exactly one debit field is provided + * + * @example + * ```typescript + * import { add, mul, tag } from '@scrawn/core'; + * + * // Using direct amount (500 cents = $5.00) + * await scrawn.sdkCallEventConsumer({ + * userId: 'user_abc123', + * debitAmount: 500 + * }); + * + * // Using price tag + * await scrawn.sdkCallEventConsumer({ + * userId: 'user_abc123', + * debitTag: 'PREMIUM_FEATURE' + * }); + * + * // Using pricing expression: (PREMIUM_CALL * 3) + EXTRA_FEE + 250 cents + * await scrawn.sdkCallEventConsumer({ + * userId: 'user_abc123', + * debitExpr: add(mul(tag('PREMIUM_CALL'), 3), tag('EXTRA_FEE'), 250) + * }); + * ``` + */ async sdkCallEventConsumer(payload: EventPayload): Promise { const validationResult = EventPayloadSchema.safeParse(payload); if (!validationResult.success) { @@ -399,15 +409,22 @@ export class Scrawn { ? SDKCallType.RAW : SDKCallType.MIDDLEWARE_CALL; - try { - log.info(`Ingesting event (type: ${eventType}) with creds: ${JSON.stringify(creds)}, payload: ${JSON.stringify(payload)}`); - - // Build debit field based on whether amount or tag is provided - const debitField = payload.debitAmount !== undefined - ? { case: 'amount' as const, value: payload.debitAmount } - : { case: 'tag' as const, value: payload.debitTag! }; - - const response = await this.grpcClient + try { + log.info(`Ingesting event (type: ${eventType}) with creds: ${JSON.stringify(creds)}, payload: ${JSON.stringify(payload)}`); + + // Build debit field based on which debit option is provided + let debitField: { case: 'amount'; value: number } | { case: 'tag'; value: string } | { case: 'expr'; value: string }; + + if (payload.debitAmount !== undefined) { + debitField = { case: 'amount' as const, value: payload.debitAmount }; + } else if (payload.debitTag !== undefined) { + debitField = { case: 'tag' as const, value: payload.debitTag }; + } else { + // debitExpr is defined (validated by schema) + debitField = { case: 'expr' as const, value: serializeExpr(payload.debitExpr!) }; + } + + const response = await this.grpcClient .newCall(EventService, 'registerEvent') .addHeader('Authorization', `Bearer ${creds.apiKey}`) .addPayload({ @@ -626,15 +643,25 @@ if (auth.postRun) await auth.postRun(); const validated = validationResult.data; - // Build input debit field - const inputDebit = validated.inputDebit.amount !== undefined - ? { case: 'inputAmount' as const, value: validated.inputDebit.amount } - : { case: 'inputTag' as const, value: validated.inputDebit.tag! }; + // Build input debit field (amount, tag, or expr) + let inputDebit: { case: 'inputAmount'; value: number } | { case: 'inputTag'; value: string } | { case: 'inputExpr'; value: string }; + if (validated.inputDebit.amount !== undefined) { + inputDebit = { case: 'inputAmount' as const, value: validated.inputDebit.amount }; + } else if (validated.inputDebit.tag !== undefined) { + inputDebit = { case: 'inputTag' as const, value: validated.inputDebit.tag }; + } else { + inputDebit = { case: 'inputExpr' as const, value: serializeExpr(validated.inputDebit.expr!) }; + } - // Build output debit field - const outputDebit = validated.outputDebit.amount !== undefined - ? { case: 'outputAmount' as const, value: validated.outputDebit.amount } - : { case: 'outputTag' as const, value: validated.outputDebit.tag! }; + // Build output debit field (amount, tag, or expr) + let outputDebit: { case: 'outputAmount'; value: number } | { case: 'outputTag'; value: string } | { case: 'outputExpr'; value: string }; + if (validated.outputDebit.amount !== undefined) { + outputDebit = { case: 'outputAmount' as const, value: validated.outputDebit.amount }; + } else if (validated.outputDebit.tag !== undefined) { + outputDebit = { case: 'outputTag' as const, value: validated.outputDebit.tag }; + } else { + outputDebit = { case: 'outputExpr' as const, value: serializeExpr(validated.outputDebit.expr!) }; + } yield { type: EventType.AI_TOKEN_USAGE, diff --git a/packages/scrawn/src/core/types/event.ts b/packages/scrawn/src/core/types/event.ts index bf7c3ca..4cbb6cd 100644 --- a/packages/scrawn/src/core/types/event.ts +++ b/packages/scrawn/src/core/types/event.ts @@ -1,50 +1,89 @@ -import { z } from 'zod'; - -/** - * Zod schema for event payload validation. - * - * Used by all event consumer methods to ensure consistent validation. - * - * Validates: - * - userId: non-empty string - * - Either debitAmount (number) OR debitTag (string), but not both - */ -export const EventPayloadSchema = z.object({ - userId: z.string().min(1, 'userId must be a non-empty string'), - debitAmount: z.number().positive('debitAmount must be a positive number').optional(), - debitTag: z.string().min(1, 'debitTag must be a non-empty string').optional(), -}).refine( - (data) => (data.debitAmount !== undefined) !== (data.debitTag !== undefined), - { message: 'Exactly one of debitAmount or debitTag must be provided' } -); +import { z } from 'zod'; +import type { PriceExpr } from '../pricing/types.js'; +import { isValidExpr } from '../pricing/validate.js'; + +/** + * Custom zod schema for PriceExpr validation. + * Validates that the value is a valid pricing expression AST. + */ +const PriceExprSchema = z.custom( + (val): val is PriceExpr => { + if (val === null || val === undefined || typeof val !== 'object') { + return false; + } + const expr = val as PriceExpr; + // Check that it has a valid kind + if (expr.kind !== 'amount' && expr.kind !== 'tag' && expr.kind !== 'op') { + return false; + } + // Use the validation function + return isValidExpr(expr); + }, + { message: 'Must be a valid pricing expression (use tag(), add(), sub(), mul(), div(), or amount())' } +); + +/** + * Zod schema for event payload validation. + * + * Used by all event consumer methods to ensure consistent validation. + * + * Validates: + * - userId: non-empty string + * - Exactly one of: debitAmount (number), debitTag (string), or debitExpr (PriceExpr) + */ +export const EventPayloadSchema = z.object({ + userId: z.string().min(1, 'userId must be a non-empty string'), + debitAmount: z.number().positive('debitAmount must be a positive number').optional(), + debitTag: z.string().min(1, 'debitTag must be a non-empty string').optional(), + debitExpr: PriceExprSchema.optional(), +}).refine( + (data) => { + const defined = [ + data.debitAmount !== undefined, + data.debitTag !== undefined, + data.debitExpr !== undefined, + ].filter(Boolean).length; + return defined === 1; + }, + { message: 'Exactly one of debitAmount, debitTag, or debitExpr must be provided' } +); -/** - * Payload structure for event tracking. - * - * Used by both sdkCallEventConsumer and middlewareEventConsumer. - * - * @property userId - The user ID associated with this event - * @property debitAmount - (Optional) Direct amount to debit for billing tracking - * @property debitTag - (Optional) Named price tag to look up amount from backend - * - * Note: Exactly one of debitAmount or debitTag must be provided. - * - * @example - * ```typescript - * // Using direct amount - * const payload1: EventPayload = { - * userId: 'u123', - * debitAmount: 5 - * }; - * - * // Using price tag - * const payload2: EventPayload = { - * userId: 'u123', - * debitTag: 'PREMIUM_FEATURE' - * }; - * ``` - */ -export type EventPayload = z.infer; +/** + * Payload structure for event tracking. + * + * Used by both sdkCallEventConsumer and middlewareEventConsumer. + * + * @property userId - The user ID associated with this event + * @property debitAmount - (Optional) Direct amount to debit in cents + * @property debitTag - (Optional) Named price tag to look up amount from backend + * @property debitExpr - (Optional) Pricing expression for complex calculations + * + * Note: Exactly one of debitAmount, debitTag, or debitExpr must be provided. + * + * @example + * ```typescript + * import { add, mul, tag } from '@scrawn/core'; + * + * // Using direct amount + * const payload1: EventPayload = { + * userId: 'u123', + * debitAmount: 500 // 500 cents = $5.00 + * }; + * + * // Using price tag + * const payload2: EventPayload = { + * userId: 'u123', + * debitTag: 'PREMIUM_FEATURE' + * }; + * + * // Using pricing expression + * const payload3: EventPayload = { + * userId: 'u123', + * debitExpr: add(mul(tag('PREMIUM_CALL'), 3), tag('EXTRA_FEE'), 250) + * }; + * ``` + */ +export type EventPayload = z.infer; /** * Generic request object type for middleware compatibility. @@ -164,15 +203,23 @@ export interface MiddlewareEventConfig { /** * Debit field schema for AI token usage. * - * Represents either a direct amount or a named price tag for billing. - * Exactly one of amount or tag must be provided. + * Represents a direct amount, a named price tag, or a pricing expression for billing. + * Exactly one of amount, tag, or expr must be provided. */ const DebitFieldSchema = z.object({ amount: z.number().nonnegative('amount must be non-negative').optional(), tag: z.string().min(1, 'tag must be a non-empty string').optional(), + expr: PriceExprSchema.optional(), }).refine( - (data) => (data.amount !== undefined) !== (data.tag !== undefined), - { message: 'Exactly one of amount or tag must be provided' } + (data) => { + const defined = [ + data.amount !== undefined, + data.tag !== undefined, + data.expr !== undefined, + ].filter(Boolean).length; + return defined === 1; + }, + { message: 'Exactly one of amount, tag, or expr must be provided' } ); /** @@ -185,8 +232,8 @@ const DebitFieldSchema = z.object({ * - model: non-empty string (e.g., 'gpt-4', 'claude-3') * - inputTokens: non-negative integer * - outputTokens: non-negative integer - * - inputDebit: either amount (number) OR tag (string), but not both - * - outputDebit: either amount (number) OR tag (string), but not both + * - inputDebit: exactly one of amount (number), tag (string), or expr (PriceExpr) + * - outputDebit: exactly one of amount (number), tag (string), or expr (PriceExpr) */ export const AITokenUsagePayloadSchema = z.object({ userId: z.string().min(1, 'userId must be a non-empty string'), @@ -207,19 +254,21 @@ export const AITokenUsagePayloadSchema = z.object({ * @property model - The AI model identifier (e.g., 'gpt-4', 'claude-3-opus') * @property inputTokens - Number of input/prompt tokens consumed * @property outputTokens - Number of output/completion tokens consumed - * @property inputDebit - Billing info for input tokens (amount or tag) - * @property outputDebit - Billing info for output tokens (amount or tag) + * @property inputDebit - Billing info for input tokens (amount, tag, or expr) + * @property outputDebit - Billing info for output tokens (amount, tag, or expr) * * @example * ```typescript + * import { mul, tag } from '@scrawn/core'; + * * // Using direct amounts * const payload1: AITokenUsagePayload = { * userId: 'u123', * model: 'gpt-4', * inputTokens: 100, * outputTokens: 50, - * inputDebit: { amount: 0.003 }, - * outputDebit: { amount: 0.006 } + * inputDebit: { amount: 3 }, // 3 cents + * outputDebit: { amount: 6 } // 6 cents * }; * * // Using price tags @@ -231,6 +280,16 @@ export const AITokenUsagePayloadSchema = z.object({ * inputDebit: { tag: 'CLAUDE_INPUT' }, * outputDebit: { tag: 'CLAUDE_OUTPUT' } * }; + * + * // Using pricing expressions (e.g., per-token pricing) + * const payload3: AITokenUsagePayload = { + * userId: 'u123', + * model: 'gpt-4', + * inputTokens: 100, + * outputTokens: 50, + * inputDebit: { expr: mul(tag('GPT4_INPUT_RATE'), 100) }, // rate * tokens + * outputDebit: { expr: mul(tag('GPT4_OUTPUT_RATE'), 50) } // rate * tokens + * }; * ``` */ export type AITokenUsagePayload = z.infer; diff --git a/packages/scrawn/src/gen/event/v1/event_pb.ts b/packages/scrawn/src/gen/event/v1/event_pb.ts index 4e31d58..cf9f799 100644 --- a/packages/scrawn/src/gen/event/v1/event_pb.ts +++ b/packages/scrawn/src/gen/event/v1/event_pb.ts @@ -137,6 +137,14 @@ export class SDKCall extends Message { */ value: string; case: "tag"; + } | { + /** + * Pricing expression (e.g., "add(mul(tag('PREMIUM'),3),250)") + * + * @generated from field: string expr = 4; + */ + value: string; + case: "expr"; } | { case: undefined; value?: undefined } = { case: undefined }; constructor(data?: PartialMessage) { @@ -150,6 +158,7 @@ export class SDKCall extends Message { { no: 1, name: "sdkCallType", kind: "enum", T: proto3.getEnumType(SDKCallType) }, { no: 2, name: "amount", kind: "scalar", T: 2 /* ScalarType.FLOAT */, oneof: "debit" }, { no: 3, name: "tag", kind: "scalar", T: 9 /* ScalarType.STRING */, oneof: "debit" }, + { no: 4, name: "expr", kind: "scalar", T: 9 /* ScalarType.STRING */, oneof: "debit" }, ]); static fromBinary(bytes: Uint8Array, options?: Partial): SDKCall { @@ -302,6 +311,14 @@ export class AITokenUsage extends Message { */ value: string; case: "inputTag"; + } | { + /** + * Pricing expression for input tokens + * + * @generated from field: string inputExpr = 8; + */ + value: string; + case: "inputExpr"; } | { case: undefined; value?: undefined } = { case: undefined }; /** @@ -319,6 +336,14 @@ export class AITokenUsage extends Message { */ value: string; case: "outputTag"; + } | { + /** + * Pricing expression for output tokens + * + * @generated from field: string outputExpr = 9; + */ + value: string; + case: "outputExpr"; } | { case: undefined; value?: undefined } = { case: undefined }; constructor(data?: PartialMessage) { @@ -334,8 +359,10 @@ export class AITokenUsage extends Message { { no: 3, name: "outputTokens", kind: "scalar", T: 5 /* ScalarType.INT32 */ }, { no: 4, name: "inputAmount", kind: "scalar", T: 2 /* ScalarType.FLOAT */, oneof: "inputDebit" }, { no: 5, name: "inputTag", kind: "scalar", T: 9 /* ScalarType.STRING */, oneof: "inputDebit" }, + { no: 8, name: "inputExpr", kind: "scalar", T: 9 /* ScalarType.STRING */, oneof: "inputDebit" }, { no: 6, name: "outputAmount", kind: "scalar", T: 2 /* ScalarType.FLOAT */, oneof: "outputDebit" }, { no: 7, name: "outputTag", kind: "scalar", T: 9 /* ScalarType.STRING */, oneof: "outputDebit" }, + { no: 9, name: "outputExpr", kind: "scalar", T: 9 /* ScalarType.STRING */, oneof: "outputDebit" }, ]); static fromBinary(bytes: Uint8Array, options?: Partial): AITokenUsage { diff --git a/packages/scrawn/src/index.ts b/packages/scrawn/src/index.ts index fdae8ae..c0b554c 100644 --- a/packages/scrawn/src/index.ts +++ b/packages/scrawn/src/index.ts @@ -1,21 +1,24 @@ -export * from './core/scrawn.js'; -export * from './core/types/event.js'; -export * from './core/types/auth.js'; - -// Export error classes for user error handling -export * from './core/errors/index.js'; - -// Export gRPC client abstraction layer -export * from './core/grpc/index.js'; - -// Export utilities -export { matchPath } from './utils/pathMatcher.js'; - -// Export generated types for advanced usage -export * from './gen/event/v1/event_connect.js'; -export * from './gen/event/v1/event_pb.js'; -export * from './gen/auth/v1/auth_connect.js'; -export * from './gen/auth/v1/auth_pb.js'; - -// Export central configuration +export * from './core/scrawn.js'; +export * from './core/types/event.js'; +export * from './core/types/auth.js'; + +// Export error classes for user error handling +export * from './core/errors/index.js'; + +// Export gRPC client abstraction layer +export * from './core/grpc/index.js'; + +// Export pricing DSL for building complex billing expressions +export * from './core/pricing/index.js'; + +// Export utilities +export { matchPath } from './utils/pathMatcher.js'; + +// Export generated types for advanced usage +export * from './gen/event/v1/event_connect.js'; +export * from './gen/event/v1/event_pb.js'; +export * from './gen/auth/v1/auth_connect.js'; +export * from './gen/auth/v1/auth_pb.js'; + +// Export central configuration export { ScrawnConfig } from './config.js'; \ No newline at end of file diff --git a/packages/scrawn/tests/unit/pricing/pricing.test.ts b/packages/scrawn/tests/unit/pricing/pricing.test.ts new file mode 100644 index 0000000..c0a4cba --- /dev/null +++ b/packages/scrawn/tests/unit/pricing/pricing.test.ts @@ -0,0 +1,279 @@ +import { describe, expect, it } from "vitest"; +import { + tag, + add, + sub, + mul, + div, + amount, + serializeExpr, + prettyPrintExpr, + validateExpr, + isValidExpr, + PricingExpressionError, +} from "../../../src/core/pricing/index.js"; +import type { PriceExpr } from "../../../src/core/pricing/types.js"; + +describe("Pricing DSL Builders", () => { + describe("tag()", () => { + it("creates a tag expression with valid name", () => { + const expr = tag("PREMIUM_CALL"); + expect(expr).toEqual({ kind: "tag", name: "PREMIUM_CALL" }); + }); + + it("accepts underscores and hyphens", () => { + expect(tag("PREMIUM_CALL")).toEqual({ kind: "tag", name: "PREMIUM_CALL" }); + expect(tag("premium-call")).toEqual({ kind: "tag", name: "premium-call" }); + expect(tag("_private")).toEqual({ kind: "tag", name: "_private" }); + }); + + it("throws on empty tag name", () => { + expect(() => tag("")).toThrow(PricingExpressionError); + expect(() => tag("")).toThrow("Tag name cannot be empty"); + }); + + it("throws on whitespace-only tag name", () => { + expect(() => tag(" ")).toThrow(PricingExpressionError); + expect(() => tag(" ")).toThrow("Tag name cannot have leading or trailing whitespace"); + }); + + it("throws on tag name with leading/trailing whitespace", () => { + expect(() => tag(" PREMIUM")).toThrow(PricingExpressionError); + expect(() => tag("PREMIUM ")).toThrow(PricingExpressionError); + }); + + it("throws on invalid tag name characters", () => { + expect(() => tag("123START")).toThrow(PricingExpressionError); + expect(() => tag("has spaces")).toThrow(PricingExpressionError); + expect(() => tag("has.dots")).toThrow(PricingExpressionError); + }); + }); + + describe("amount()", () => { + it("creates an amount expression", () => { + expect(amount(100)).toEqual({ kind: "amount", value: 100 }); + expect(amount(0)).toEqual({ kind: "amount", value: 0 }); + expect(amount(-50)).toEqual({ kind: "amount", value: -50 }); + }); + + it("throws on non-integer values", () => { + expect(() => amount(2.5)).toThrow(PricingExpressionError); + expect(() => amount(99.99)).toThrow("Amount must be an integer (cents)"); + }); + + it("throws on non-finite values", () => { + expect(() => amount(Infinity)).toThrow(PricingExpressionError); + expect(() => amount(-Infinity)).toThrow(PricingExpressionError); + expect(() => amount(NaN)).toThrow(PricingExpressionError); + }); + }); + + describe("add()", () => { + it("creates an addition expression", () => { + const expr = add(100, 200); + expect(expr.kind).toBe("op"); + expect(expr.op).toBe("ADD"); + expect(expr.args).toHaveLength(2); + }); + + it("auto-wraps numbers as amounts", () => { + const expr = add(100, tag("FEE"), 50); + expect(expr.args[0]).toEqual({ kind: "amount", value: 100 }); + expect(expr.args[1]).toEqual({ kind: "tag", name: "FEE" }); + expect(expr.args[2]).toEqual({ kind: "amount", value: 50 }); + }); + + it("throws on fewer than 2 arguments", () => { + expect(() => add(100 as any)).toThrow(PricingExpressionError); + expect(() => add(100 as any)).toThrow("requires at least 2 arguments"); + }); + + it("supports many arguments", () => { + const expr = add(1, 2, 3, 4, 5); + expect(expr.args).toHaveLength(5); + }); + }); + + describe("sub()", () => { + it("creates a subtraction expression", () => { + const expr = sub(100, 50); + expect(expr.op).toBe("SUB"); + }); + + it("throws on fewer than 2 arguments", () => { + expect(() => sub(100 as any)).toThrow(PricingExpressionError); + }); + }); + + describe("mul()", () => { + it("creates a multiplication expression", () => { + const expr = mul(tag("PER_TOKEN"), 100); + expect(expr.op).toBe("MUL"); + }); + + it("throws on fewer than 2 arguments", () => { + expect(() => mul(100 as any)).toThrow(PricingExpressionError); + }); + }); + + describe("div()", () => { + it("creates a division expression", () => { + const expr = div(tag("TOTAL"), 2); + expect(expr.op).toBe("DIV"); + }); + + it("throws on literal division by zero", () => { + expect(() => div(100, 0)).toThrow(PricingExpressionError); + expect(() => div(100, 0)).toThrow("Division by zero"); + }); + + it("allows division by tag (backend validates at runtime)", () => { + // This should NOT throw - backend will handle tag resolution + const expr = div(100, tag("DIVISOR")); + expect(expr.op).toBe("DIV"); + }); + + it("detects division by zero in any divisor position", () => { + expect(() => div(100, 2, 0)).toThrow("Division by zero"); + expect(() => div(100, 0, 2)).toThrow("Division by zero"); + }); + }); + + describe("nested expressions", () => { + it("supports deeply nested expressions", () => { + const expr = add(mul(tag("PREMIUM_CALL"), 3), tag("EXTRA_FEE"), 250); + expect(expr.kind).toBe("op"); + expect(expr.op).toBe("ADD"); + expect(expr.args).toHaveLength(3); + + const mulExpr = expr.args[0]; + expect(mulExpr.kind).toBe("op"); + if (mulExpr.kind === "op") { + expect(mulExpr.op).toBe("MUL"); + } + }); + + it("validates nested expressions", () => { + expect(() => add(mul(tag(""), 3), 100)).toThrow(PricingExpressionError); + }); + }); +}); + +describe("Pricing DSL Serialization", () => { + describe("serializeExpr()", () => { + it("serializes amount expressions", () => { + expect(serializeExpr(amount(100))).toBe("100"); + expect(serializeExpr(amount(0))).toBe("0"); + expect(serializeExpr(amount(-50))).toBe("-50"); + }); + + it("serializes tag expressions", () => { + expect(serializeExpr(tag("PREMIUM"))).toBe("tag('PREMIUM')"); + expect(serializeExpr(tag("API_CALL"))).toBe("tag('API_CALL')"); + }); + + it("serializes simple operations", () => { + expect(serializeExpr(add(100, 200))).toBe("add(100,200)"); + expect(serializeExpr(sub(100, 50))).toBe("sub(100,50)"); + expect(serializeExpr(mul(10, 5))).toBe("mul(10,5)"); + expect(serializeExpr(div(100, 2))).toBe("div(100,2)"); + }); + + it("serializes mixed operations", () => { + expect(serializeExpr(add(100, tag("FEE")))).toBe("add(100,tag('FEE'))"); + }); + + it("serializes nested operations", () => { + const expr = add(mul(tag("PREMIUM_CALL"), 3), tag("EXTRA_FEE"), 250); + expect(serializeExpr(expr)).toBe( + "add(mul(tag('PREMIUM_CALL'),3),tag('EXTRA_FEE'),250)" + ); + }); + + it("serializes complex expressions", () => { + const expr = div(add(mul(tag("INPUT"), 2), mul(tag("OUTPUT"), 3)), 100); + expect(serializeExpr(expr)).toBe( + "div(add(mul(tag('INPUT'),2),mul(tag('OUTPUT'),3)),100)" + ); + }); + }); + + describe("prettyPrintExpr()", () => { + it("pretty prints simple expressions", () => { + expect(prettyPrintExpr(amount(100))).toBe("100"); + expect(prettyPrintExpr(tag("PREMIUM"))).toBe("tag('PREMIUM')"); + }); + + it("pretty prints operations with indentation", () => { + const expr = add(100, 200); + const output = prettyPrintExpr(expr); + expect(output).toContain("add("); + expect(output).toContain("100"); + expect(output).toContain("200"); + }); + }); +}); + +describe("Pricing DSL Validation", () => { + describe("validateExpr()", () => { + it("validates amount expressions", () => { + expect(() => validateExpr(amount(100))).not.toThrow(); + expect(() => validateExpr({ kind: "amount", value: 2.5 } as PriceExpr)).toThrow(); + }); + + it("validates tag expressions", () => { + expect(() => validateExpr(tag("VALID"))).not.toThrow(); + expect(() => validateExpr({ kind: "tag", name: "" } as PriceExpr)).toThrow(); + }); + + it("validates operation expressions", () => { + expect(() => validateExpr(add(100, 200))).not.toThrow(); + expect(() => + validateExpr({ kind: "op", op: "ADD", args: [amount(100)] } as PriceExpr) + ).toThrow(); + }); + }); + + describe("isValidExpr()", () => { + it("returns true for valid expressions", () => { + expect(isValidExpr(amount(100))).toBe(true); + expect(isValidExpr(tag("PREMIUM"))).toBe(true); + expect(isValidExpr(add(100, 200))).toBe(true); + }); + + it("returns false for invalid expressions", () => { + expect(isValidExpr({ kind: "amount", value: 2.5 } as PriceExpr)).toBe(false); + expect(isValidExpr({ kind: "tag", name: "" } as PriceExpr)).toBe(false); + }); + }); +}); + +describe("Integration Examples", () => { + it("handles typical billing expression", () => { + // (PREMIUM_CALL * 3) + EXTRA_FEE + 250 cents + const expr = add(mul(tag("PREMIUM_CALL"), 3), tag("EXTRA_FEE"), 250); + const serialized = serializeExpr(expr); + expect(serialized).toBe("add(mul(tag('PREMIUM_CALL'),3),tag('EXTRA_FEE'),250)"); + }); + + it("handles per-token pricing", () => { + // (INPUT_TOKENS * INPUT_RATE) + (OUTPUT_TOKENS * OUTPUT_RATE) + const expr = add( + mul(tag("INPUT_TOKENS"), tag("INPUT_RATE")), + mul(tag("OUTPUT_TOKENS"), tag("OUTPUT_RATE")) + ); + const serialized = serializeExpr(expr); + expect(serialized).toBe( + "add(mul(tag('INPUT_TOKENS'),tag('INPUT_RATE')),mul(tag('OUTPUT_TOKENS'),tag('OUTPUT_RATE')))" + ); + }); + + it("handles discount calculation", () => { + // SUBTOTAL - (SUBTOTAL * DISCOUNT_PERCENT / 100) + const expr = sub(tag("SUBTOTAL"), div(mul(tag("SUBTOTAL"), tag("DISCOUNT_PERCENT")), 100)); + const serialized = serializeExpr(expr); + expect(serialized).toBe( + "sub(tag('SUBTOTAL'),div(mul(tag('SUBTOTAL'),tag('DISCOUNT_PERCENT')),100))" + ); + }); +}); diff --git a/packages/scrawn/tests/unit/types/aiTokenUsagePayload.test.ts b/packages/scrawn/tests/unit/types/aiTokenUsagePayload.test.ts index fc61584..5c33114 100644 --- a/packages/scrawn/tests/unit/types/aiTokenUsagePayload.test.ts +++ b/packages/scrawn/tests/unit/types/aiTokenUsagePayload.test.ts @@ -1,5 +1,6 @@ import { describe, expect, it } from "vitest"; import { AITokenUsagePayloadSchema } from "../../../src/core/types/event.js"; +import { mul, tag, add } from "../../../src/core/pricing/index.js"; describe("AITokenUsagePayloadSchema", () => { describe("valid payloads", () => { @@ -29,6 +30,32 @@ describe("AITokenUsagePayloadSchema", () => { expect(result.success).toBe(true); }); + it("accepts payloads with expr-based debits", () => { + const result = AITokenUsagePayloadSchema.safeParse({ + userId: "user_1", + model: "gpt-4", + inputTokens: 100, + outputTokens: 50, + inputDebit: { expr: mul(tag("GPT4_INPUT_RATE"), 100) }, + outputDebit: { expr: mul(tag("GPT4_OUTPUT_RATE"), 50) }, + }); + + expect(result.success).toBe(true); + }); + + it("accepts payloads with complex expr-based debits", () => { + const result = AITokenUsagePayloadSchema.safeParse({ + userId: "user_1", + model: "gpt-4", + inputTokens: 100, + outputTokens: 50, + inputDebit: { expr: add(mul(tag("BASE_RATE"), 100), tag("PREMIUM_FEE")) }, + outputDebit: { expr: mul(tag("OUTPUT_RATE"), 50) }, + }); + + expect(result.success).toBe(true); + }); + it("accepts payloads with mixed debit types", () => { const result = AITokenUsagePayloadSchema.safeParse({ userId: "user_1", @@ -42,6 +69,32 @@ describe("AITokenUsagePayloadSchema", () => { expect(result.success).toBe(true); }); + it("accepts payloads mixing expr with amount", () => { + const result = AITokenUsagePayloadSchema.safeParse({ + userId: "user_1", + model: "gpt-4", + inputTokens: 100, + outputTokens: 50, + inputDebit: { expr: mul(tag("INPUT_RATE"), 100) }, + outputDebit: { amount: 6 }, + }); + + expect(result.success).toBe(true); + }); + + it("accepts payloads mixing expr with tag", () => { + const result = AITokenUsagePayloadSchema.safeParse({ + userId: "user_1", + model: "gpt-4", + inputTokens: 100, + outputTokens: 50, + inputDebit: { tag: "INPUT_TAG" }, + outputDebit: { expr: mul(tag("OUTPUT_RATE"), 50) }, + }); + + expect(result.success).toBe(true); + }); + it("accepts payloads with zero tokens", () => { const result = AITokenUsagePayloadSchema.safeParse({ userId: "user_1", @@ -148,6 +201,45 @@ describe("AITokenUsagePayloadSchema", () => { expect(result.success).toBe(false); }); + it("rejects payloads with both amount and expr in inputDebit", () => { + const result = AITokenUsagePayloadSchema.safeParse({ + userId: "user_1", + model: "gpt-4", + inputTokens: 100, + outputTokens: 50, + inputDebit: { amount: 3, expr: tag("INPUT") }, + outputDebit: { amount: 6 }, + }); + + expect(result.success).toBe(false); + }); + + it("rejects payloads with both tag and expr in outputDebit", () => { + const result = AITokenUsagePayloadSchema.safeParse({ + userId: "user_1", + model: "gpt-4", + inputTokens: 100, + outputTokens: 50, + inputDebit: { amount: 3 }, + outputDebit: { tag: "OUTPUT", expr: tag("OUTPUT_EXPR") }, + }); + + expect(result.success).toBe(false); + }); + + it("rejects payloads with all three in debit", () => { + const result = AITokenUsagePayloadSchema.safeParse({ + userId: "user_1", + model: "gpt-4", + inputTokens: 100, + outputTokens: 50, + inputDebit: { amount: 3, tag: "TAG", expr: tag("EXPR") }, + outputDebit: { amount: 6 }, + }); + + expect(result.success).toBe(false); + }); + it("rejects payloads with empty inputDebit", () => { const result = AITokenUsagePayloadSchema.safeParse({ userId: "user_1", @@ -200,6 +292,19 @@ describe("AITokenUsagePayloadSchema", () => { expect(result.success).toBe(false); }); + it("rejects payloads with invalid expr", () => { + const result = AITokenUsagePayloadSchema.safeParse({ + userId: "user_1", + model: "gpt-4", + inputTokens: 100, + outputTokens: 50, + inputDebit: { expr: { invalid: "expression" } }, + outputDebit: { amount: 6 }, + }); + + expect(result.success).toBe(false); + }); + it("rejects payloads missing required fields", () => { const result = AITokenUsagePayloadSchema.safeParse({ userId: "user_1", diff --git a/packages/scrawn/tests/unit/types/eventPayload.test.ts b/packages/scrawn/tests/unit/types/eventPayload.test.ts index c29dec0..808a96f 100644 --- a/packages/scrawn/tests/unit/types/eventPayload.test.ts +++ b/packages/scrawn/tests/unit/types/eventPayload.test.ts @@ -1,5 +1,6 @@ import { describe, expect, it } from "vitest"; import { EventPayloadSchema } from "../../../src/core/types/event.js"; +import { add, mul, tag } from "../../../src/core/pricing/index.js"; describe("EventPayloadSchema", () => { it("accepts payloads with debitAmount", () => { @@ -20,7 +21,25 @@ describe("EventPayloadSchema", () => { expect(result.success).toBe(true); }); - it("rejects payloads with both debit fields", () => { + it("accepts payloads with debitExpr (simple tag)", () => { + const result = EventPayloadSchema.safeParse({ + userId: "user_1", + debitExpr: tag("PREMIUM_CALL"), + }); + + expect(result.success).toBe(true); + }); + + it("accepts payloads with debitExpr (complex expression)", () => { + const result = EventPayloadSchema.safeParse({ + userId: "user_1", + debitExpr: add(mul(tag("PREMIUM_CALL"), 3), tag("EXTRA_FEE"), 250), + }); + + expect(result.success).toBe(true); + }); + + it("rejects payloads with both debitAmount and debitTag", () => { const result = EventPayloadSchema.safeParse({ userId: "user_1", debitAmount: 5, @@ -30,6 +49,37 @@ describe("EventPayloadSchema", () => { expect(result.success).toBe(false); }); + it("rejects payloads with both debitAmount and debitExpr", () => { + const result = EventPayloadSchema.safeParse({ + userId: "user_1", + debitAmount: 5, + debitExpr: tag("PREMIUM"), + }); + + expect(result.success).toBe(false); + }); + + it("rejects payloads with both debitTag and debitExpr", () => { + const result = EventPayloadSchema.safeParse({ + userId: "user_1", + debitTag: "PREMIUM", + debitExpr: tag("OTHER"), + }); + + expect(result.success).toBe(false); + }); + + it("rejects payloads with all three debit fields", () => { + const result = EventPayloadSchema.safeParse({ + userId: "user_1", + debitAmount: 5, + debitTag: "PREMIUM", + debitExpr: tag("OTHER"), + }); + + expect(result.success).toBe(false); + }); + it("rejects payloads without debit info", () => { const result = EventPayloadSchema.safeParse({ userId: "user_1", @@ -46,4 +96,23 @@ describe("EventPayloadSchema", () => { expect(result.success).toBe(false); }); + + it("rejects invalid debitExpr (not a valid PriceExpr)", () => { + const result = EventPayloadSchema.safeParse({ + userId: "user_1", + debitExpr: { invalid: "expression" }, + }); + + expect(result.success).toBe(false); + }); + + it("rejects debitExpr with invalid nested expression", () => { + // Manually construct an invalid expression (non-integer amount) + const result = EventPayloadSchema.safeParse({ + userId: "user_1", + debitExpr: { kind: "amount", value: 2.5 }, + }); + + expect(result.success).toBe(false); + }); });