From 996d1d42fbfb670ac6746292f81e1e2437f583a8 Mon Sep 17 00:00:00 2001 From: Devyash Saini Date: Tue, 27 Jan 2026 12:18:46 +0530 Subject: [PATCH 1/3] refactor(grpc): shared grpc request builder context to remove redundency Signed-off-by: Devyash Saini --- packages/scrawn/src/core/grpc/callContext.ts | 113 ++++++ packages/scrawn/src/core/grpc/client.ts | 309 +++++++-------- packages/scrawn/src/core/grpc/index.ts | 11 +- .../scrawn/src/core/grpc/requestBuilder.ts | 352 +++++++++--------- .../src/core/grpc/streamRequestBuilder.ts | 62 ++- packages/scrawn/src/core/grpc/types.ts | 47 +++ .../tests/unit/grpc/requestBuilder.test.ts | 25 +- 7 files changed, 543 insertions(+), 376 deletions(-) create mode 100644 packages/scrawn/src/core/grpc/callContext.ts diff --git a/packages/scrawn/src/core/grpc/callContext.ts b/packages/scrawn/src/core/grpc/callContext.ts new file mode 100644 index 0000000..43fda73 --- /dev/null +++ b/packages/scrawn/src/core/grpc/callContext.ts @@ -0,0 +1,113 @@ +/** + * Shared call context for gRPC request builders. + * + * This module encapsulates the common state and operations used by both + * `RequestBuilder` (unary) and `StreamRequestBuilder` (client-streaming), + * eliminating code duplication while keeping each builder focused on its + * specific execution model. + * + * Design: + * - The context is immutable after construction; headers are accumulated via + * a mutable map owned by the context but exposed through a controlled API. + * - Builders receive the context via constructor injection and delegate header + * management and logging to it. + * - The context does NOT own payload or execution logic; those remain in the + * individual builders to preserve separation of concerns. + */ + +import type { Client, Transport } from '@connectrpc/connect'; +import type { ServiceType } from '@bufbuild/protobuf'; +import { createClient } from '@connectrpc/connect'; +import { ScrawnLogger } from '../../utils/logger.js'; +import type { ServiceMethodNames, Headers } from './types.js'; + +/** + * Shared context for a single gRPC call. + * + * Holds the typed client, method name, headers, and logger. Provides helper + * methods used by both unary and streaming request builders. + * + * @template S - The gRPC service type + * @template M - The method name on the service (string literal) + */ +export class GrpcCallContext< + S extends ServiceType, + M extends ServiceMethodNames +> { + /** The typed Connect client for the service */ + public readonly client: Client; + + /** The method name being invoked */ + public readonly methodName: M; + + /** Accumulated headers for the call */ + private readonly headers: Headers = {}; + + /** Logger scoped to the builder type */ + public readonly log: ScrawnLogger; + + /** + * Constructs a new GrpcCallContext. + * + * @param transport - The Connect transport to use + * @param service - The gRPC service definition + * @param methodName - The method to invoke + * @param loggerName - Name for the logger (e.g., 'RequestBuilder') + */ + constructor( + transport: Transport, + service: S, + methodName: M, + loggerName: string + ) { + this.client = createClient(service, transport); + this.methodName = methodName; + this.log = new ScrawnLogger(loggerName); + } + + /** + * Add a header to the call. + * + * @param key - Header name + * @param value - Header value + */ + addHeader(key: string, value: string): void { + this.headers[key] = value; + } + + /** + * Get a shallow copy of the current headers. + * Used when making the actual gRPC call. + */ + getHeaders(): Headers { + return { ...this.headers }; + } + + /** + * Log the start of a call (info level). + */ + logCallStart(): void { + this.log.info(`Making gRPC call to ${String(this.methodName)}`); + this.log.debug(`Headers: ${JSON.stringify(this.headers)}`); + } + + /** + * Log successful completion of a call (info level). + */ + logCallSuccess(): void { + this.log.info(`Successfully completed gRPC call to ${String(this.methodName)}`); + } + + /** + * Log a call failure (error level). + * + * @param error - The error that occurred + */ + logCallError(error: unknown): void { + this.log.error( + `gRPC call to ${String(this.methodName)} failed: ${ + error instanceof Error ? error.message : 'Unknown error' + }` + ); + } +} diff --git a/packages/scrawn/src/core/grpc/client.ts b/packages/scrawn/src/core/grpc/client.ts index ea72df9..22ea648 100644 --- a/packages/scrawn/src/core/grpc/client.ts +++ b/packages/scrawn/src/core/grpc/client.ts @@ -1,156 +1,163 @@ -/** - * Type-safe gRPC client abstraction layer. - * - * This module provides the main entry point for making gRPC calls with - * full compile-time type safety and a beautiful fluent API. - * - * @example - * ```typescript - * const client = new GrpcClient('https://api.scrawn.dev'); - * - * const response = await client - * .newCall(EventService, 'registerEvent') - * .addHeader('Authorization', `Bearer ${apiKey}`) - * .addHeader('x-request-id', '123') - * .addPayload({ - * type: EventType.SDK_CALL, - * userId: 'user_123', - * data: { case: 'sdkCall', value: new SDKCall({ sdkCallType: SDKCallType.RAW, debitAmount: 10 }) } - * }) - * .request(); - * ``` - */ - +/** + * Type-safe gRPC client abstraction layer. + * + * This module provides the main entry point for making gRPC calls with + * full compile-time type safety and a beautiful fluent API. + * + * @example + * ```typescript + * const client = new GrpcClient('https://api.scrawn.dev'); + * + * const response = await client + * .newCall(EventService, 'registerEvent') + * .addHeader('Authorization', `Bearer ${apiKey}`) + * .addHeader('x-request-id', '123') + * .addPayload({ + * type: EventType.SDK_CALL, + * userId: 'user_123', + * data: { case: 'sdkCall', value: new SDKCall({ sdkCallType: SDKCallType.RAW, debitAmount: 10 }) } + * }) + * .request(); + * ``` + */ + import type { Transport } from '@connectrpc/connect'; import type { ServiceType } from '@bufbuild/protobuf'; import { createConnectTransport } from '@connectrpc/connect-node'; +import { GrpcCallContext } from './callContext.js'; import { RequestBuilder } from './requestBuilder.js'; import { StreamRequestBuilder } from './streamRequestBuilder.js'; import { ScrawnLogger } from '../../utils/logger.js'; import type { ServiceMethodNames } from './types.js'; import { ScrawnConfig } from '../../config.js'; - -const log = new ScrawnLogger('GrpcClient'); - -/** - * Main gRPC client for making type-safe API calls. - * - * This class manages the underlying transport and provides a factory method - * for creating type-safe request builders. Each request builder is bound to - * a specific service and method, ensuring compile-time type safety. - * - * The client handles: - * - Transport configuration (HTTP/1.1) - * - Request building with fluent API - * - Header management (auth should be added via .addHeader()) - * - Error handling and logging - * - * @example - * ```typescript - * // Initialize the client - * const client = new GrpcClient('https://api.scrawn.dev'); - * - * // Fetch credentials from your auth system - * const creds = await getCredentials(); - * - * // Make a call to EventService.registerEvent - * const eventResponse = await client - * .newCall(EventService, 'registerEvent') - * .addHeader('Authorization', `Bearer ${creds.apiKey}`) - * .addPayload({ type: EventType.SDK_CALL, userId: 'u123', ... }) - * .request(); - * - * // Make a call to AuthService.createAPIKey - * const authResponse = await client - * .newCall(AuthService, 'createAPIKey') - * .addHeader('Authorization', `Bearer ${creds.apiKey}`) - * .addPayload({ name: 'My Key', expiresIn: 3600n }) - * .request(); - * ``` - */ -export class GrpcClient { - /** The underlying gRPC transport */ - private readonly transport: Transport; - - /** Base URL for API calls */ - private readonly baseURL: string; - - /** - * Create a new GrpcClient. - * - * @param baseURL - The base URL of the gRPC API (e.g., 'https://api.scrawn.dev') - * - * @example - * ```typescript - * const client = new GrpcClient('https://api.scrawn.dev'); - * ``` - */ - constructor(baseURL: string) { - this.baseURL = baseURL; - - log.info(`Initializing gRPC client for ${baseURL}`); - - this.transport = createConnectTransport({ - baseUrl: this.baseURL, - httpVersion: ScrawnConfig.grpc.httpVersion, - useBinaryFormat: true, // Use binary protobuf (smaller than JSON) - }); - - log.info(`gRPC client initialized with HTTP/${ScrawnConfig.grpc.httpVersion}`); - } - - /** - * Create a new request builder for a specific service method. - * - * This is the entry point for making gRPC calls. The method is fully type-safe: - * - Service parameter must be a valid gRPC service - * - Method name must exist on the service (autocomplete provided) - * - Payload type is inferred from the method - * - Response type is inferred from the method - * - * @template S - The gRPC service type (auto-inferred) - * @template M - The method name (auto-inferred and validated) - * - * @param service - The gRPC service definition (e.g., EventService, AuthService) - * @param method - The method name as a string literal (e.g., 'registerEvent', 'createAPIKey') - * @returns A new RequestBuilder for chaining headers, payload, and execution - * - * @example - * ```typescript - * // EventService.registerEvent - * const eventBuilder = client.newCall(EventService, 'registerEvent'); - * // Payload type is RegisterEventRequest - * // Response type is RegisterEventResponse - * - * // AuthService.createAPIKey - * const authBuilder = client.newCall(AuthService, 'createAPIKey'); - * // Payload type is CreateAPIKeyRequest - * // Response type is CreateAPIKeyResponse - * - * // Type error - method doesn't exist! - * // const invalid = client.newCall(EventService, 'nonExistentMethod'); - * ``` - */ - newCall>( - service: S, - method: M - ): RequestBuilder { - log.debug(`Creating new request builder for ${service.typeName}.${method}`); - return new RequestBuilder(this.transport, service, method); - } - - /** - * Get the base URL of this client. - * - * @returns The base URL string - */ - getBaseURL(): string { - return this.baseURL; - } - + +const log = new ScrawnLogger('GrpcClient'); + /** + * Main gRPC client for making type-safe API calls. + * + * This class manages the underlying transport and provides factory methods + * for creating type-safe request builders. Each request builder is bound to + * a specific service and method, ensuring compile-time type safety. + * + * The client handles: + * - Transport configuration (HTTP/1.1) + * - Request building with fluent API + * - Header management (auth should be added via .addHeader()) + * - Error handling and logging + * + * @example + * ```typescript + * // Initialize the client + * const client = new GrpcClient('https://api.scrawn.dev'); + * + * // Fetch credentials from your auth system + * const creds = await getCredentials(); + * + * // Make a call to EventService.registerEvent + * const eventResponse = await client + * .newCall(EventService, 'registerEvent') + * .addHeader('Authorization', `Bearer ${creds.apiKey}`) + * .addPayload({ type: EventType.SDK_CALL, userId: 'u123', ... }) + * .request(); + * + * // Make a call to AuthService.createAPIKey + * const authResponse = await client + * .newCall(AuthService, 'createAPIKey') + * .addHeader('Authorization', `Bearer ${creds.apiKey}`) + * .addPayload({ name: 'My Key', expiresIn: 3600n }) + * .request(); + * ``` + */ +export class GrpcClient { + /** The underlying gRPC transport */ + private readonly transport: Transport; + + /** Base URL for API calls */ + private readonly baseURL: string; + + /** + * Create a new GrpcClient. + * + * @param baseURL - The base URL of the gRPC API (e.g., 'https://api.scrawn.dev') + * + * @example + * ```typescript + * const client = new GrpcClient('https://api.scrawn.dev'); + * ``` + */ + constructor(baseURL: string) { + this.baseURL = baseURL; + + log.info(`Initializing gRPC client for ${baseURL}`); + + this.transport = createConnectTransport({ + baseUrl: this.baseURL, + httpVersion: ScrawnConfig.grpc.httpVersion, + useBinaryFormat: true, // Use binary protobuf (smaller than JSON) + }); + + log.info(`gRPC client initialized with HTTP/${ScrawnConfig.grpc.httpVersion}`); + } + + /** + * Create a new request builder for a specific service method. + * + * This is the entry point for making unary gRPC calls. The method is fully type-safe: + * - Service parameter must be a valid gRPC service + * - Method name must exist on the service (autocomplete provided) + * - Payload type is inferred from the method + * - Response type is inferred from the method + * + * @template S - The gRPC service type (auto-inferred) + * @template M - The method name (auto-inferred and validated) + * + * @param service - The gRPC service definition (e.g., EventService, AuthService) + * @param method - The method name as a string literal (e.g., 'registerEvent', 'createAPIKey') + * @returns A new RequestBuilder for chaining headers, payload, and execution + * + * @example + * ```typescript + * // EventService.registerEvent + * const eventBuilder = client.newCall(EventService, 'registerEvent'); + * // Payload type is RegisterEventRequest + * // Response type is RegisterEventResponse + * + * // AuthService.createAPIKey + * const authBuilder = client.newCall(AuthService, 'createAPIKey'); + * // Payload type is CreateAPIKeyRequest + * // Response type is CreateAPIKeyResponse + * + * // Type error - method doesn't exist! + * // const invalid = client.newCall(EventService, 'nonExistentMethod'); + * ``` + */ + newCall>( + service: S, + method: M + ): RequestBuilder { + log.debug(`Creating new request builder for ${service.typeName}.${String(method)}`); + const ctx = new GrpcCallContext( + this.transport, + service, + method, + 'RequestBuilder' + ); + return new RequestBuilder(ctx); + } + + /** + * Get the base URL of this client. + * + * @returns The base URL string + */ + getBaseURL(): string { + return this.baseURL; + } + + /** * Get the underlying transport (for advanced use cases). - * + * * @returns The Connect transport instance * @internal */ @@ -160,20 +167,20 @@ export class GrpcClient { /** * Create a new stream request builder for a client-streaming service method. - * + * * This is the entry point for making client-streaming gRPC calls. The method is fully type-safe: * - Service parameter must be a valid gRPC service * - Method name must exist on the service (autocomplete provided) * - Payload type is inferred from the method * - Response type is inferred from the method - * + * * @template S - The gRPC service type (auto-inferred) * @template M - The method name (auto-inferred and validated) - * + * * @param service - The gRPC service definition (e.g., EventService) * @param method - The method name as a string literal (e.g., 'streamEvents') * @returns A new StreamRequestBuilder for chaining headers and streaming payloads - * + * * @example * ```typescript * // EventService.streamEvents (client-streaming) @@ -189,6 +196,12 @@ export class GrpcClient { method: M ): StreamRequestBuilder { log.debug(`Creating new stream request builder for ${service.typeName}.${String(method)}`); - return new StreamRequestBuilder(this.transport, service, method); + const ctx = new GrpcCallContext( + this.transport, + service, + method, + 'StreamRequestBuilder' + ); + return new StreamRequestBuilder(ctx); } } diff --git a/packages/scrawn/src/core/grpc/index.ts b/packages/scrawn/src/core/grpc/index.ts index d82f377..930cdc4 100644 --- a/packages/scrawn/src/core/grpc/index.ts +++ b/packages/scrawn/src/core/grpc/index.ts @@ -1,13 +1,14 @@ /** * gRPC abstraction layer - Type-safe fluent API for gRPC calls. - * + * * This module provides a beautiful, type-safe interface for making gRPC calls * with automatic type inference, compile-time validation, and a fluent API. - * + * * @module grpc */ export { GrpcClient } from './client.js'; +export { GrpcCallContext } from './callContext.js'; export { RequestBuilder } from './requestBuilder.js'; export { StreamRequestBuilder } from './streamRequestBuilder.js'; export type { @@ -16,4 +17,10 @@ export type { MethodOutput, Headers, RequestState, + MethodKind, + MethodsOfKind, + UnaryMethodNames, + ClientStreamingMethodNames, + ServerStreamingMethodNames, + BidiStreamingMethodNames, } from './types.js'; diff --git a/packages/scrawn/src/core/grpc/requestBuilder.ts b/packages/scrawn/src/core/grpc/requestBuilder.ts index 1663b31..2e1be7a 100644 --- a/packages/scrawn/src/core/grpc/requestBuilder.ts +++ b/packages/scrawn/src/core/grpc/requestBuilder.ts @@ -1,184 +1,168 @@ -/** - * Type-safe fluent API for building and executing gRPC requests. - * - * This module provides a beautiful chain-able interface that ensures: - * - Compile-time type safety for all operations - * - Autocomplete for service methods and payload fields - * - Runtime validation of request state - * - Clean separation of concerns - * - * @example - * ```typescript - * const response = await client - * .newCall(EventService, 'registerEvent') - * .addHeader('x-request-id', '123') - * .addPayload({ - * type: EventType.SDK_CALL, - * userId: 'user_123', - * data: { case: 'sdkCall', value: { sdkCallType: SDKCallType.RAW, debitAmount: 10 } } - * }) - * .request(); - * ``` - */ - -import type { Client, Transport } from '@connectrpc/connect'; -import type { ServiceType } from '@bufbuild/protobuf'; -import { createClient } from '@connectrpc/connect'; -import { ScrawnLogger } from '../../utils/logger.js'; -import type { - ServiceMethodNames, - MethodInput, - MethodOutput, - Headers, - RequestState, -} from './types.js'; - -const log = new ScrawnLogger('RequestBuilder'); - -/** - * Builder for constructing type-safe gRPC requests. - * - * This class implements a fluent interface where each method returns `this`, - * allowing for method chaining. Type parameters ensure that the payload type - * matches the selected service method at compile time. - * - * @template S - The gRPC service type - * @template M - The method name on the service (string literal) - */ -export class RequestBuilder< - S extends ServiceType, - M extends ServiceMethodNames -> { - private readonly client: Client; - private readonly methodName: M; - private readonly headers: Headers = {}; - private payload: any = null; - private readonly state: RequestState = { hasPayload: false }; - - /** - * @internal - * Constructs a new RequestBuilder. Use GrpcClient.newCall() instead. - */ - constructor( - transport: Transport, - service: S, - methodName: M - ) { - // Create a typed client for this service - this.client = createClient(service, transport); - this.methodName = methodName; - } - - /** - * Add a header to the request. - * - * Headers are sent with the gRPC call and can be used for authentication, - * tracing, or custom metadata. This method can be called multiple times - * to add multiple headers. - * - * @param key - Header name (e.g., 'authorization', 'x-request-id') - * @param value - Header value - * @returns The builder instance for chaining - * - * @example - * ```typescript - * builder - * .addHeader('authorization', 'Bearer token123') - * .addHeader('x-request-id', 'abc-def-123') - * ``` - */ - addHeader(key: string, value: string): this { - this.headers[key] = value; - return this; - } - - /** - * Set the request payload. - * - * The payload type is automatically inferred from the service method, - * providing full autocomplete and type checking. This method can only - * be called once per request to prevent accidental overwrites. - * - * @param payload - The request payload (type-checked against the method's input type) - * @returns The builder instance for chaining - * @throws Error if payload has already been set - * - * @example - * ```typescript - * builder.addPayload({ - * type: EventType.SDK_CALL, - * userId: 'user_123', - * data: { - * case: 'sdkCall', - * value: { - * sdkCallType: SDKCallType.RAW, - * debitAmount: 10 - * } - * } - * }) - * ``` - */ - addPayload(payload: MethodInput extends infer T ? Partial : never): this { - if (this.state.hasPayload) { - throw new Error( - 'Payload has already been set. Cannot add payload multiple times. ' + - 'Create a new request builder if you need to make another call.' - ); - } - - this.payload = payload; - this.state.hasPayload = true; - return this; - } - - /** - * Execute the gRPC request. - * - * Validates that a payload has been set, then makes the actual gRPC call - * using the configured service, method, headers, and payload. The response - * type is automatically inferred from the service method. - * - * @returns A promise that resolves to the typed response - * @throws Error if no payload has been set - * @throws Error if the gRPC call fails - * - * @example - * ```typescript - * const response = await builder - * .addPayload({ userId: 'user_123', ... }) - * .request(); - * - * // response is fully typed based on the method's output type - * console.log(response.random); - * ``` - */ - async request(): Promise> { - if (!this.state.hasPayload || this.payload === null) { - throw new Error( - 'Cannot make request without payload. Call addPayload() first.' - ); - } - - try { - log.info(`Making gRPC call to ${this.methodName}`); - log.debug(`Headers: ${JSON.stringify(this.headers)}`); - log.debug(`Payload: ${JSON.stringify(this.payload)}`); - - // The actual gRPC call - fully type-safe! - const response = await (this.client[this.methodName] as any)( - this.payload, - { headers: this.headers } - ); - - log.info(`Successfully completed gRPC call to ${this.methodName}`); - return response as MethodOutput; - } catch (error) { - log.error( - `gRPC call to ${this.methodName} failed: ${ - error instanceof Error ? error.message : 'Unknown error' - }` - ); - throw error; - } - } -} +/** + * Type-safe fluent API for building and executing unary gRPC requests. + * + * This module provides a chain-able interface that ensures: + * - Compile-time type safety for all operations + * - Autocomplete for service methods and payload fields + * - Runtime validation of request state + * - Clean separation of concerns + * + * @example + * ```typescript + * const response = await client + * .newCall(EventService, 'registerEvent') + * .addHeader('x-request-id', '123') + * .addPayload({ + * type: EventType.SDK_CALL, + * userId: 'user_123', + * data: { case: 'sdkCall', value: { sdkCallType: SDKCallType.RAW, debitAmount: 10 } } + * }) + * .request(); + * ``` + */ + +import type { ServiceType } from '@bufbuild/protobuf'; +import type { + ServiceMethodNames, + MethodInput, + MethodOutput, + RequestState, +} from './types.js'; +import type { GrpcCallContext } from './callContext.js'; + +/** + * Builder for constructing type-safe unary gRPC requests. + * + * This class implements a fluent interface where each method returns `this`, + * allowing for method chaining. Type parameters ensure that the payload type + * matches the selected service method at compile time. + * + * @template S - The gRPC service type + * @template M - The method name on the service (string literal) + */ +export class RequestBuilder< + S extends ServiceType, + M extends ServiceMethodNames +> { + private readonly ctx: GrpcCallContext; + private payload: MethodInput | null = null; + private readonly state: RequestState = { hasPayload: false }; + + /** + * @internal + * Constructs a new RequestBuilder. Use GrpcClient.newCall() instead. + * + * @param ctx - The shared call context (injected by GrpcClient) + */ + constructor(ctx: GrpcCallContext) { + this.ctx = ctx; + } + + /** + * Add a header to the request. + * + * Headers are sent with the gRPC call and can be used for authentication, + * tracing, or custom metadata. This method can be called multiple times + * to add multiple headers. + * + * @param key - Header name (e.g., 'authorization', 'x-request-id') + * @param value - Header value + * @returns The builder instance for chaining + * + * @example + * ```typescript + * builder + * .addHeader('authorization', 'Bearer token123') + * .addHeader('x-request-id', 'abc-def-123') + * ``` + */ + addHeader(key: string, value: string): this { + this.ctx.addHeader(key, value); + return this; + } + + /** + * Set the request payload. + * + * The payload type is automatically inferred from the service method, + * providing full autocomplete and type checking. This method can only + * be called once per request to prevent accidental overwrites. + * + * @param payload - The request payload (type-checked against the method's input type) + * @returns The builder instance for chaining + * @throws Error if payload has already been set + * + * @example + * ```typescript + * builder.addPayload({ + * type: EventType.SDK_CALL, + * userId: 'user_123', + * data: { + * case: 'sdkCall', + * value: { + * sdkCallType: SDKCallType.RAW, + * debitAmount: 10 + * } + * } + * }) + * ``` + */ + addPayload(payload: MethodInput extends infer T ? Partial : never): this { + if (this.state.hasPayload) { + throw new Error( + 'Payload has already been set. Cannot add payload multiple times. ' + + 'Create a new request builder if you need to make another call.' + ); + } + + this.payload = payload as unknown as MethodInput; + this.state.hasPayload = true; + return this; + } + + /** + * Execute the gRPC request. + * + * Validates that a payload has been set, then makes the actual gRPC call + * using the configured service, method, headers, and payload. The response + * type is automatically inferred from the service method. + * + * @returns A promise that resolves to the typed response + * @throws Error if no payload has been set + * @throws Error if the gRPC call fails + * + * @example + * ```typescript + * const response = await builder + * .addPayload({ userId: 'user_123', ... }) + * .request(); + * + * // response is fully typed based on the method's output type + * console.log(response.random); + * ``` + */ + async request(): Promise> { + if (!this.state.hasPayload || this.payload === null) { + throw new Error( + 'Cannot make request without payload. Call addPayload() first.' + ); + } + + try { + 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() } + ); + + this.ctx.logCallSuccess(); + return response as MethodOutput; + } 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 c5d6d47..6988de0 100644 --- a/packages/scrawn/src/core/grpc/streamRequestBuilder.ts +++ b/packages/scrawn/src/core/grpc/streamRequestBuilder.ts @@ -1,12 +1,12 @@ /** * Type-safe fluent API for building and executing client-streaming gRPC requests. - * + * * This module provides a chain-able interface for client-streaming RPC calls that ensures: * - Compile-time type safety for all operations * - Autocomplete for service methods and payload fields * - Runtime validation of request state * - Non-blocking stream consumption with internal buffering - * + * * @example * ```typescript * const response = await client @@ -16,26 +16,21 @@ * ``` */ -import type { Client, Transport } from '@connectrpc/connect'; import type { ServiceType } from '@bufbuild/protobuf'; -import { createClient } from '@connectrpc/connect'; -import { ScrawnLogger } from '../../utils/logger.js'; import type { ServiceMethodNames, MethodInput, MethodOutput, - Headers, } from './types.js'; - -const log = new ScrawnLogger('StreamRequestBuilder'); +import type { GrpcCallContext } from './callContext.js'; /** * Builder for constructing type-safe client-streaming gRPC requests. - * + * * This class implements a fluent interface where each method returns `this`, * allowing for method chaining. Type parameters ensure that the payload type * matches the selected service method at compile time. - * + * * @template S - The gRPC service type * @template M - The method name on the service (string literal) */ @@ -43,35 +38,29 @@ export class StreamRequestBuilder< S extends ServiceType, M extends ServiceMethodNames > { - private readonly client: Client; - private readonly methodName: M; - private readonly headers: Headers = {}; + private readonly ctx: GrpcCallContext; /** * @internal * Constructs a new StreamRequestBuilder. Use GrpcClient.newStreamCall() instead. + * + * @param ctx - The shared call context (injected by GrpcClient) */ - constructor( - transport: Transport, - service: S, - methodName: M - ) { - // Create a typed client for this service - this.client = createClient(service, transport); - this.methodName = methodName; + constructor(ctx: GrpcCallContext) { + this.ctx = ctx; } /** * Add a header to the request. - * + * * Headers are sent with the gRPC call and can be used for authentication, * tracing, or custom metadata. This method can be called multiple times * to add multiple headers. - * + * * @param key - Header name (e.g., 'authorization', 'x-request-id') * @param value - Header value * @returns The builder instance for chaining - * + * * @example * ```typescript * builder @@ -80,28 +69,28 @@ export class StreamRequestBuilder< * ``` */ addHeader(key: string, value: string): this { - this.headers[key] = value; + this.ctx.addHeader(key, value); return this; } /** * Execute the client-streaming gRPC request with the provided async iterable. - * + * * Consumes the async iterable and streams each item to the server. * The iterable is consumed in the background to avoid blocking the caller. * Returns a promise that resolves when the stream is complete and the server responds. - * + * * @param iterable - An async iterable of request payloads to stream * @returns A promise that resolves to the typed response * @throws Error if the gRPC call fails - * + * * @example * ```typescript * async function* generatePayloads() { * yield { type: EventType.AI_TOKEN_USAGE, userId: 'u123', data: { ... } }; * yield { type: EventType.AI_TOKEN_USAGE, userId: 'u123', data: { ... } }; * } - * + * * const response = await builder.stream(generatePayloads()); * console.log(`Processed ${response.eventsProcessed} events`); * ``` @@ -110,23 +99,18 @@ export class StreamRequestBuilder< iterable: AsyncIterable extends infer T ? Partial : never> ): Promise> { try { - log.info(`Starting client-streaming gRPC call to ${String(this.methodName)}`); - log.debug(`Headers: ${JSON.stringify(this.headers)}`); + this.ctx.logCallStart(); // The actual client-streaming gRPC call - const response = await (this.client[this.methodName])( + const response = await (this.ctx.client[this.ctx.methodName] as any)( iterable, - { headers: this.headers } + { headers: this.ctx.getHeaders() } ); - log.info(`Successfully completed streaming gRPC call to ${String(this.methodName)}`); + this.ctx.logCallSuccess(); return response as MethodOutput; } catch (error) { - log.error( - `Streaming gRPC call to ${String(this.methodName)} failed: ${ - error instanceof Error ? error.message : 'Unknown 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 56e87b8..8dad68d 100644 --- a/packages/scrawn/src/core/grpc/types.ts +++ b/packages/scrawn/src/core/grpc/types.ts @@ -68,3 +68,50 @@ export interface Headers { export interface RequestState { hasPayload: boolean; } + +/** + * Represents the kind of a gRPC method. + * Connect RPC method info includes `kind` with these values. + */ +export type MethodKind = 'unary' | 'server_streaming' | 'client_streaming' | 'bidi_streaming'; + +/** + * Extract methods of a specific kind from a service. + * + * @template S - The gRPC service type + * @template K - The method kind to filter by + */ +export type MethodsOfKind< + S extends ServiceType, + K extends MethodKind +> = { + [M in keyof S['methods'] & string]: S['methods'][M] extends { kind: infer MK } + ? MK extends K + ? M + : never + : never; +}[keyof S['methods'] & string]; + +/** + * Extract unary method names from a service. + * Unary methods accept a single request and return a single response. + */ +export type UnaryMethodNames = MethodsOfKind; + +/** + * Extract client-streaming method names from a service. + * Client-streaming methods accept a stream of requests and return a single response. + */ +export type ClientStreamingMethodNames = MethodsOfKind; + +/** + * Extract server-streaming method names from a service. + * Server-streaming methods accept a single request and return a stream of responses. + */ +export type ServerStreamingMethodNames = MethodsOfKind; + +/** + * Extract bidirectional-streaming method names from a service. + * Bidi-streaming methods accept a stream of requests and return a stream of responses. + */ +export type BidiStreamingMethodNames = MethodsOfKind; diff --git a/packages/scrawn/tests/unit/grpc/requestBuilder.test.ts b/packages/scrawn/tests/unit/grpc/requestBuilder.test.ts index ac2a243..8fc146b 100644 --- a/packages/scrawn/tests/unit/grpc/requestBuilder.test.ts +++ b/packages/scrawn/tests/unit/grpc/requestBuilder.test.ts @@ -1,5 +1,6 @@ import { describe, expect, it } from "vitest"; import { RequestBuilder } from "../../../src/core/grpc/requestBuilder.js"; +import { GrpcCallContext } from "../../../src/core/grpc/callContext.js"; import { createMockTransport } from "../../mocks/mockTransport.js"; import { PaymentService } from "../../../src/gen/payment/v1/payment_connect.js"; import { CreateCheckoutLinkResponse } from "../../../src/gen/payment/v1/payment_pb.js"; @@ -16,7 +17,13 @@ describe("RequestBuilder", () => { }, }); - const builder = new RequestBuilder(transport, PaymentService, "createCheckoutLink"); + const ctx = new GrpcCallContext( + transport, + PaymentService, + "createCheckoutLink", + "RequestBuilder" + ); + const builder = new RequestBuilder(ctx); const response = await builder .addHeader("Authorization", "Bearer token") .addPayload({ userId: "user_1" }) @@ -30,7 +37,13 @@ describe("RequestBuilder", () => { unary: () => new CreateCheckoutLinkResponse({ checkoutLink: "" }), }); - const builder = new RequestBuilder(transport, PaymentService, "createCheckoutLink"); + const ctx = new GrpcCallContext( + transport, + PaymentService, + "createCheckoutLink", + "RequestBuilder" + ); + const builder = new RequestBuilder(ctx); await expect(builder.request()).rejects.toThrow("addPayload"); }); @@ -39,7 +52,13 @@ describe("RequestBuilder", () => { unary: () => new CreateCheckoutLinkResponse({ checkoutLink: "" }), }); - const builder = new RequestBuilder(transport, PaymentService, "createCheckoutLink"); + const ctx = new GrpcCallContext( + transport, + PaymentService, + "createCheckoutLink", + "RequestBuilder" + ); + const builder = new RequestBuilder(ctx); builder.addPayload({ userId: "user_1" }); expect(() => builder.addPayload({ userId: "user_2" })).toThrow( From f66fd675dd6d18540eefefc2ddbbe8336405d556 Mon Sep 17 00:00:00 2001 From: Devyash Saini Date: Tue, 27 Jan 2026 13:14:51 +0530 Subject: [PATCH 2/3] refactor(grpc): replace any in request builders Signed-off-by: Devyash Saini --- .../scrawn/src/core/grpc/requestBuilder.ts | 19 ++++-- .../src/core/grpc/streamRequestBuilder.ts | 18 +++-- packages/scrawn/src/core/grpc/types.ts | 68 +++++++++++++++++-- 3 files changed, 90 insertions(+), 15 deletions(-) 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..bfbe0da 100644 --- a/packages/scrawn/src/core/grpc/types.ts +++ b/packages/scrawn/src/core/grpc/types.ts @@ -110,8 +110,66 @@ export type ClientStreamingMethodNames = MethodsOfKind = MethodsOfKind; -/** - * Extract bidirectional-streaming method names from a service. - * Bidi-streaming methods accept a stream of requests and return a stream of responses. - */ -export type BidiStreamingMethodNames = MethodsOfKind; +/** + * Extract bidirectional-streaming method names from a service. + * Bidi-streaming methods accept a stream of requests and return a stream of responses. + */ +export type BidiStreamingMethodNames = 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; From 608141c42dea731206be120528279c7d067b5374 Mon Sep 17 00:00:00 2001 From: Devyash Saini Date: Tue, 27 Jan 2026 13:17:41 +0530 Subject: [PATCH 3/3] feat(tags): arithematic operations for dynamic pricing factors Signed-off-by: Devyash Saini --- packages/scrawn/proto/event/v1/event.proto | 3 + packages/scrawn/src/core/pricing/builders.ts | 167 +++++++++++ packages/scrawn/src/core/pricing/index.ts | 42 +++ packages/scrawn/src/core/pricing/serialize.ts | 123 ++++++++ packages/scrawn/src/core/pricing/types.ts | 56 ++++ packages/scrawn/src/core/pricing/validate.ts | 151 ++++++++++ packages/scrawn/src/core/scrawn.ts | 117 +++++--- packages/scrawn/src/core/types/event.ts | 171 +++++++---- packages/scrawn/src/gen/event/v1/event_pb.ts | 27 ++ packages/scrawn/src/index.ts | 43 +-- .../scrawn/tests/unit/pricing/pricing.test.ts | 279 ++++++++++++++++++ .../unit/types/aiTokenUsagePayload.test.ts | 105 +++++++ .../tests/unit/types/eventPayload.test.ts | 71 ++++- 13 files changed, 1233 insertions(+), 122 deletions(-) create mode 100644 packages/scrawn/src/core/pricing/builders.ts create mode 100644 packages/scrawn/src/core/pricing/index.ts create mode 100644 packages/scrawn/src/core/pricing/serialize.ts create mode 100644 packages/scrawn/src/core/pricing/types.ts create mode 100644 packages/scrawn/src/core/pricing/validate.ts create mode 100644 packages/scrawn/tests/unit/pricing/pricing.test.ts 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/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); + }); });