From 69651e02b965eb8f5258c5b514df05745e5ef3e2 Mon Sep 17 00:00:00 2001 From: Ben Papillon Date: Wed, 3 Jun 2026 14:20:38 -0700 Subject: [PATCH] add credit lease + reservation SDK with optional Redis backend Client-side credit lease + reservation gating on check/trackWithReservation, with in-memory or Redis-backed lease and reservation stores. Adds prewarm, getCreditBalance (lease-aware, expired-lease safe), and lease-matched credit selection on mixed-credit flags. --- src/cache/redis.ts | 31 +- src/credits/check.ts | 397 +++++++++++++ src/credits/index.ts | 23 + src/credits/lease-manager.ts | 195 +++++++ src/credits/lease-store.ts | 186 ++++++ src/credits/redis-lease-store.ts | 223 +++++++ src/credits/redis-reservation-store.ts | 277 +++++++++ src/credits/reservation-store.ts | 137 +++++ src/credits/track.ts | 62 ++ src/credits/types.ts | 176 ++++++ src/datastream/datastream-client.ts | 38 +- src/index.ts | 8 + src/rules-engine.ts | 31 +- src/wrapper.ts | 443 +++++++++++++- tests/unit/credits/check-and-track.test.ts | 544 ++++++++++++++++++ tests/unit/credits/fake-redis.ts | 240 ++++++++ tests/unit/credits/lease-manager.test.ts | 383 ++++++++++++ tests/unit/credits/lease-store.test.ts | 172 ++++++ tests/unit/credits/redis-lease-store.test.ts | 161 ++++++ .../credits/redis-reservation-store.test.ts | 246 ++++++++ tests/unit/credits/reservation-store.test.ts | 107 ++++ tests/unit/credits/wasm-credit-gate.test.ts | 227 ++++++++ .../unit/datastream/datastream-client.test.ts | 16 +- tests/unit/wrapper.test.ts | 240 +++++++- 24 files changed, 4543 insertions(+), 20 deletions(-) create mode 100644 src/credits/check.ts create mode 100644 src/credits/index.ts create mode 100644 src/credits/lease-manager.ts create mode 100644 src/credits/lease-store.ts create mode 100644 src/credits/redis-lease-store.ts create mode 100644 src/credits/redis-reservation-store.ts create mode 100644 src/credits/reservation-store.ts create mode 100644 src/credits/track.ts create mode 100644 src/credits/types.ts create mode 100644 tests/unit/credits/check-and-track.test.ts create mode 100644 tests/unit/credits/fake-redis.ts create mode 100644 tests/unit/credits/lease-manager.test.ts create mode 100644 tests/unit/credits/lease-store.test.ts create mode 100644 tests/unit/credits/redis-lease-store.test.ts create mode 100644 tests/unit/credits/redis-reservation-store.test.ts create mode 100644 tests/unit/credits/reservation-store.test.ts create mode 100644 tests/unit/credits/wasm-credit-gate.test.ts diff --git a/src/cache/redis.ts b/src/cache/redis.ts index 474623af..5e65142a 100644 --- a/src/cache/redis.ts +++ b/src/cache/redis.ts @@ -1,8 +1,11 @@ import { CacheProvider, CacheOptions } from "./types"; /** - * Minimal interface describing the redis client methods used by RedisCacheProvider. - * Compatible with the 'redis' package's RedisClientType. + * Minimal interface describing the redis client methods used by the SDK. + * Compatible with the 'redis' package's RedisClientType (node-redis v4). + * + * Includes hash + sorted-set + eval surface used by the credit-lease and + * reservation stores to coordinate state across multiple SDK pods. */ export interface RedisClient { get(key: string): Promise; @@ -10,6 +13,30 @@ export interface RedisClient { setEx(key: string, seconds: number, value: string): Promise; del(key: string | string[]): Promise; scanIterator(options: { MATCH: string; COUNT: number }): AsyncIterable; + // Hash ops — used to store lease + reservation state as a single field-set + // so partial updates (e.g. atomic decrement on localRemainingCredits) don't + // step on neighboring fields. + hSet(key: string, field: string | Record, value?: string | number): Promise; + hGet(key: string, field: string): Promise; + hGetAll(key: string): Promise>; + hDel(key: string, field: string | string[]): Promise; + // Sorted-set ops — used to index reservations by expiry timestamp so the + // sweeper can pop expired entries in O(log n). + zAdd(key: string, members: { score: number; value: string } | { score: number; value: string }[]): Promise; + zRangeByScore(key: string, min: number | string, max: number | string): Promise; + zRem(key: string, member: string | string[]): Promise; + zCard(key: string): Promise; + // Lua scripts — used for the single-key atomic primitives (check-and-decrement + // on lease balance, claim-and-delete on a reservation). Kept single-key so + // they're safe under Redis Cluster (a multi-key EVAL whose keys span slots + // raises CROSSSLOT). + eval( + script: string, + options: { keys: string[]; arguments: string[] }, + ): Promise; + // Expiry on a millisecond-precision absolute timestamp — used to auto-clean + // lease + reservation rows shortly after their declared expiry. + pExpireAt(key: string, timestamp: number): Promise; } export interface RedisOptions extends CacheOptions { diff --git a/src/credits/check.ts b/src/credits/check.ts new file mode 100644 index 00000000..15c1204c --- /dev/null +++ b/src/credits/check.ts @@ -0,0 +1,397 @@ +import { randomUUID } from "crypto"; + +import type * as api from "../api"; +import type { DataStreamClient } from "../datastream"; +import type { Logger } from "../logger"; +import type { WasmFeatureEntitlement } from "../rules-engine"; +import type { CheckFlagOptions } from "../wrapper"; + +import { CreditLeaseManager } from "./lease-manager"; +import type { ILeaseStore } from "./lease-store"; +import type { IReservationStore } from "./reservation-store"; +import type { + CheckOptions, + CheckResult, + Reservation, + ResolvedLeaseConfig, +} from "./types"; + +/** Internal helper bundling everything needed to satisfy a lease-bearing check. */ +export interface CreditCheckDeps { + leaseStore: ILeaseStore; + reservations: IReservationStore; + manager: CreditLeaseManager; + datastream: DataStreamClient | undefined; + logger: Logger; +} + +/** + * Drives a single `client.check` with `usage` / `eventUsage` set. + * + * Steps: + * 1. Resolve the matching credit-balance condition on the flag to get + * `creditId` + `consumptionRate`. + * 2. Acquire (or reuse) a lease for `(company, creditId)`. + * 3. Try to reserve `quantity × consumptionRate` from the lease. + * 4. Run the WASM rules engine against a substituted company snapshot + * (`credit_balances[creditId] = lease.localRemaining` *before* the + * reservation we just made was debited), plus `event_usage` options so + * the engine evaluates the post-call balance. The reservation we made + * in step 3 only sticks if the engine says allowed. + */ +export async function checkWithLease( + deps: CreditCheckDeps, + key: string, + evalCtx: api.CheckFlagRequestBody, + options: CheckOptions, + fallback: () => Promise, +): Promise { + const { datastream, leaseStore, reservations, manager, logger } = deps; + const onFailure = options.onAcquireFailure ?? "fail-closed"; + + if (!datastream) { + logger.debug( + "Credit-lease check requested without datastream — falling back to plain check", + ); + return fallback(); + } + + let flag: api.RulesengineFlag | null = null; + try { + flag = await datastream.getFlag(key); + } catch (err) { + logger.warn(`Lease check: failed to load flag ${key}: ${err}`); + } + if (!flag) { + logger.debug(`Lease check: no cached flag for ${key}, falling back`); + return fallback(); + } + + const { eventSubtype, quantity } = extractPreflightQuantity(options); + + const company = evalCtx.company ? await datastream.getCachedCompany(evalCtx.company) : null; + const user = evalCtx.user ? await datastream.getCachedUser(evalCtx.user) : null; + if (evalCtx.company && !company) { + logger.debug(`Lease check: company not in cache for keys ${JSON.stringify(evalCtx.company)}, falling back`); + return fallback(); + } + if (!company) { + logger.debug("Lease check: no company on evalCtx, falling back"); + return fallback(); + } + + // A flag can carry credit conditions from several plans, each metering a + // different credit type. Lease the credit the company's *matched* plan + // entitlement actually uses — not whichever credit condition is declared + // first on the flag. We can't read that off the flag structurally (a + // condition doesn't know which rule the company matched), so we probe the + // engine: its entitlement reports the plan-correct creditId regardless of + // balance. Falls back to first-match for single-credit flags or when the + // probe can't resolve a credit. + const match = await resolveCreditCondition(datastream, flag, company, user, eventSubtype, logger); + if (!match) { + logger.debug( + `Lease check: flag ${key} has no matching credit condition (subtype=${eventSubtype ?? ""}), falling back`, + ); + return fallback(); + } + + const creditId = match.condition.creditId; + const consumptionRate = match.condition.consumptionRate ?? 0; + if (consumptionRate <= 0) { + logger.debug(`Lease check: condition has no consumption_rate, falling back`); + return fallback(); + } + const creditCost = quantity * consumptionRate; + + const lease = await manager.acquireIfNeeded(company.id, creditId); + if (!lease) { + return handleAcquireFailure(onFailure, key, "lease_acquire_failed", flag); + } + + const reservedLocally = await leaseStore.tryReserve(company.id, creditId, creditCost); + if (!reservedLocally) { + // Lease has less than `creditCost` left locally. Pass `creditCost` so + // `maybeExtendInBackground` extends even when the ratio is still above + // the low-watermark (e.g. a single large request). + await manager.maybeExtendInBackground(company.id, creditId, creditCost); + const retry = await leaseStore.tryReserve(company.id, creditId, creditCost); + if (!retry) { + return handleAcquireFailure(onFailure, key, "insufficient_lease_balance", flag); + } + } + + // Record the reservation *before* the WASM eval. The debit above and this + // record are two steps; persisting the hold first means a crash between + // them leaves a sweepable reservation (refunded at its TTL) rather than + // stranding the debited credits until the whole lease expires. The + // remaining unprotected window is just the gap between the debit and this + // add (no I/O in between); a crash there leaks at most `creditCost` until + // the lease's own expiry reclaims it server-side. + const reservation = registerReservation({ + leaseId: lease.leaseId, + companyId: company.id, + creditTypeId: creditId, + eventSubtype: eventSubtype ?? match.condition.eventSubtype ?? "", + quantityReserved: quantity, + creditsReserved: creditCost, + consumptionRate, + reservationTTL: resolveReservationTTL(deps, creditId), + evalCtx, + }); + await reservations.add(reservation); + + // Substitute the lease balance into the company snapshot so the engine + // gates against the lease's local view, not the server's authoritative + // balance. We use the *pre-reservation* localRemaining + event_usage so + // the engine computes `pre - qty*rate >= 0`, which matches the plan. + // Pre-reservation = current store balance (post-tryReserve) + creditCost + // we just debited. + const postReservationEntry = await leaseStore.get(company.id, creditId); + const preReservation = (postReservationEntry?.localRemainingCredits ?? 0) + creditCost; + const substituted = substituteCreditBalance(company, creditId, preReservation); + + const wasmOptions = buildWasmOptions(options); + let result; + try { + result = await datastream + .getRulesEngine() + .checkFlagWithOptions(flag, substituted, user ?? null, wasmOptions); + } catch (err) { + logger.error(`Lease check: WASM eval failed: ${err}`); + // Cancel the hold: removes the reservation record and refunds the full + // reserved amount back to the lease in one atomic step. + await reservations.consume(reservation.id, 0); + return handleAcquireFailure(onFailure, key, `wasm_error: ${err}`, flag); + } + + if (!result.value) { + await reservations.consume(reservation.id, 0); + return { + allowed: false, + value: false, + reason: result.reason ?? "denied_by_engine", + entitlement: normalizeEntitlement(result.entitlement), + flagKey: result.flagKey ?? key, + flagId: result.flagId, + }; + } + + // Fire-and-forget low-water-mark refresh now that we've debited. + void manager.maybeExtendInBackground(company.id, creditId); + + return { + allowed: true, + value: true, + reservation, + reason: result.reason ?? "lease_reserved", + entitlement: normalizeEntitlement(result.entitlement), + flagKey: result.flagKey ?? key, + flagId: result.flagId, + }; +} + +function resolveReservationTTL(deps: CreditCheckDeps, creditId: string): ResolvedLeaseConfig { + return deps.manager.resolveConfig(creditId); +} + +function extractPreflightQuantity(options: CheckOptions): { + eventSubtype: string | undefined; + quantity: number; +} { + if (options.eventUsage) { + const entries = Object.entries(options.eventUsage); + if (entries.length > 0) { + // Only one subtype per check is supported for the reservation. If + // multiple are passed, take the first; the engine will still + // gate on whichever conditions match. + const [subtype, qty] = entries[0]; + return { eventSubtype: subtype, quantity: qty }; + } + } + if (options.usage !== undefined) { + return { eventSubtype: undefined, quantity: options.usage }; + } + return { eventSubtype: undefined, quantity: 0 }; +} + +function buildWasmOptions(options: CheckOptions): CheckFlagOptions | undefined { + const out: CheckFlagOptions = {}; + if (options.eventUsage && Object.keys(options.eventUsage).length > 0) { + out.eventUsage = options.eventUsage; + } + if (options.usage !== undefined) { + out.usage = options.usage; + } + return Object.keys(out).length > 0 ? out : undefined; +} + +interface CreditConditionMatch { + condition: api.RulesengineCondition & { creditId: string }; +} + +/** + * Resolve the credit condition to lease against. A flag can declare credit + * conditions for several credit types — one per plan that entitles the feature + * — and the right one for this company is the credit its *matched* plan + * entitlement uses, which only the rules engine knows. + * + * When the flag meters this event subtype in a single credit type (the common + * case) the first matching condition is unambiguous and we return it directly, + * with no extra engine call. Only when conditions disagree on credit type do + * we probe: the engine's `entitlement.creditId` names the company's metered + * credit regardless of balance, and we select that credit's condition to read + * its `consumption_rate`. A probe that can't resolve a credit (non-credit + * entitlement or an error) falls back to first-match, preserving prior + * behavior. + */ +async function resolveCreditCondition( + datastream: DataStreamClient, + flag: api.RulesengineFlag, + company: api.RulesengineCompany, + user: object | null, + eventSubtype: string | undefined, + logger: Logger, +): Promise { + const first = findCreditCondition(flag, eventSubtype); + if (!first) return null; + if (collectCreditIds(flag, eventSubtype).size <= 1) return first; + + try { + const probe = await datastream.getRulesEngine().checkFlagWithOptions(flag, company, user, null); + const creditId = probe.entitlement?.creditId; + if (creditId) { + const byCredit = findCreditCondition(flag, eventSubtype, creditId); + if (byCredit) return byCredit; + logger.debug( + `Lease check: entitlement credit ${creditId} has no matching condition on flag, falling back to first credit condition`, + ); + } + } catch (err) { + logger.warn(`Lease check: credit-resolution probe failed (${err}), falling back to first credit condition`); + } + return first; +} + +/** Distinct credit-type ids across the flag's credit conditions for `eventSubtype`. */ +function collectCreditIds(flag: api.RulesengineFlag, eventSubtype: string | undefined): Set { + const ids = new Set(); + const scan = (conditions: api.RulesengineCondition[]) => { + for (const c of conditions) { + if (c.conditionType !== "credit" || !c.creditId) continue; + if (eventSubtype !== undefined && c.eventSubtype !== eventSubtype) continue; + ids.add(c.creditId); + } + }; + for (const rule of flag.rules ?? []) { + scan(rule.conditions ?? []); + for (const group of rule.conditionGroups ?? []) scan(group.conditions ?? []); + } + return ids; +} + +function findCreditCondition( + flag: api.RulesengineFlag, + eventSubtype: string | undefined, + creditId?: string, +): CreditConditionMatch | null { + for (const rule of flag.rules ?? []) { + const inRule = matchInConditions(rule.conditions ?? [], eventSubtype, creditId); + if (inRule) return inRule; + for (const group of rule.conditionGroups ?? []) { + const inGroup = matchInConditions(group.conditions ?? [], eventSubtype, creditId); + if (inGroup) return inGroup; + } + } + return null; +} + +function matchInConditions( + conditions: api.RulesengineCondition[], + eventSubtype: string | undefined, + creditId?: string, +): CreditConditionMatch | null { + for (const c of conditions) { + if (c.conditionType !== "credit") continue; + if (!c.creditId) continue; + if (eventSubtype !== undefined && c.eventSubtype !== eventSubtype) continue; + if (creditId !== undefined && c.creditId !== creditId) continue; + return { condition: c as api.RulesengineCondition & { creditId: string } }; + } + return null; +} + +function substituteCreditBalance( + company: api.RulesengineCompany, + creditId: string, + balance: number, +): api.RulesengineCompany { + return { + ...company, + creditBalances: { + ...(company.creditBalances ?? {}), + [creditId]: balance, + }, + }; +} + +function registerReservation(args: { + leaseId: string; + companyId: string; + creditTypeId: string; + eventSubtype: string; + quantityReserved: number; + creditsReserved: number; + consumptionRate: number; + reservationTTL: ResolvedLeaseConfig; + evalCtx: api.CheckFlagRequestBody; +}): Reservation { + return { + id: randomUUID(), + leaseId: args.leaseId, + companyId: args.companyId, + creditTypeId: args.creditTypeId, + eventSubtype: args.eventSubtype, + quantityReserved: args.quantityReserved, + creditsReserved: args.creditsReserved, + consumptionRate: args.consumptionRate, + expiresAt: new Date(Date.now() + args.reservationTTL.reservationTTL), + evalCtx: args.evalCtx, + }; +} + +function normalizeEntitlement( + raw: WasmFeatureEntitlement | undefined, +): api.RulesengineFeatureEntitlement | undefined { + if (!raw) return undefined; + return { + ...raw, + metricResetAt: raw.metricResetAt ? new Date(raw.metricResetAt) : undefined, + }; +} + +function handleAcquireFailure( + mode: "fail-open" | "fail-closed", + flagKey: string, + reason: string, + flag: api.RulesengineFlag | null, +): CheckResult { + if (mode === "fail-closed") { + return { + allowed: false, + value: false, + reason, + flagKey, + flagId: flag?.id, + err: reason, + }; + } + return { + allowed: true, + value: true, + reason: `${reason}_fail_open`, + flagKey, + flagId: flag?.id, + err: reason, + }; +} diff --git a/src/credits/index.ts b/src/credits/index.ts new file mode 100644 index 00000000..1678b3d7 --- /dev/null +++ b/src/credits/index.ts @@ -0,0 +1,23 @@ +export { LeaseStore, leaseKey, type LeaseEntry, type ILeaseStore } from "./lease-store"; +export { ReservationStore, type IReservationStore } from "./reservation-store"; +export { CreditLeaseManager } from "./lease-manager"; +export { RedisLeaseStore } from "./redis-lease-store"; +export { RedisReservationStore } from "./redis-reservation-store"; +export { + DEFAULT_LEASE_DURATION_MS, + DEFAULT_LEASE_SIZE, + DEFAULT_LOW_WATER_MARK, + DEFAULT_PREWARM_POLL_INTERVAL_MS, + DEFAULT_PREWARM_RESOLVE_TIMEOUT_MS, + DEFAULT_RESERVATION_TTL_MS, + DEFAULT_SWEEP_INTERVAL_MS, +} from "./types"; +export type { + CreditLeaseConfig, + ResolvedLeaseConfig, + Reservation, + CheckOptions, + CheckResult, + OnAcquireFailure, + TrackWithReservationOptions, +} from "./types"; diff --git a/src/credits/lease-manager.ts b/src/credits/lease-manager.ts new file mode 100644 index 00000000..03f2abd3 --- /dev/null +++ b/src/credits/lease-manager.ts @@ -0,0 +1,195 @@ +import type * as api from "../api"; +import type { CreditsClient } from "../api/resources/credits/client/Client"; +import type { Logger } from "../logger"; + +import { type ILeaseStore, type LeaseEntry, leaseKey } from "./lease-store"; +import { + DEFAULT_LEASE_DURATION_MS, + DEFAULT_LEASE_SIZE, + DEFAULT_LOW_WATER_MARK, + DEFAULT_RESERVATION_TTL_MS, + type CreditLeaseConfig, + type ResolvedLeaseConfig, +} from "./types"; + +/** + * Owns the lifecycle of `credit_lease` rows for a single client: acquire on + * first use or after expiry, extend when the local view dips below the low + * water mark, release on `client.close()`. + * + * Concurrency: each operation (acquire, extend) has its own single-flight + * map keyed by `(company, creditType)`. Concurrent callers of the *same* + * operation share one wire request; a concurrent acquire + extend on the + * same slot are two independent wire requests (they're semantically + * different and must not share state). + */ +export class CreditLeaseManager { + private readonly creditsClient: CreditsClient; + private readonly leaseStore: ILeaseStore; + private readonly logger: Logger; + private readonly config: CreditLeaseConfig; + // Per-operation single-flight maps keyed by `(company, creditType)`. Kept + // separate so a concurrent `maybeExtendInBackground` doesn't accidentally + // return the in-flight acquire promise (or vice-versa) — the shapes + // match, but the semantics don't. + private readonly inflightAcquire = new Map>(); + private readonly inflightExtend = new Map>(); + + constructor(opts: { + creditsClient: CreditsClient; + leaseStore: ILeaseStore; + logger: Logger; + config: CreditLeaseConfig; + }) { + this.creditsClient = opts.creditsClient; + this.leaseStore = opts.leaseStore; + this.logger = opts.logger; + this.config = opts.config; + } + + resolveConfig(creditTypeId: string): ResolvedLeaseConfig { + const override = this.config.overrides?.[creditTypeId]; + return { + leaseDuration: override?.defaultLeaseDuration ?? this.config.defaultLeaseDuration ?? DEFAULT_LEASE_DURATION_MS, + reservationTTL: override?.defaultReservationTTL ?? this.config.defaultReservationTTL ?? DEFAULT_RESERVATION_TTL_MS, + leaseSize: override?.defaultLeaseSize ?? this.config.defaultLeaseSize ?? DEFAULT_LEASE_SIZE, + lowWaterMark: override?.lowWaterMark ?? this.config.lowWaterMark ?? DEFAULT_LOW_WATER_MARK, + }; + } + + /** + * Return the current lease entry, acquiring one (or replacing an expired + * one) if none is live. Single-flight so racing callers share one request. + */ + async acquireIfNeeded(companyId: string, creditTypeId: string): Promise { + const existing = await this.leaseStore.get(companyId, creditTypeId); + if (existing && existing.expiresAt.getTime() > Date.now()) { + return existing; + } + // An expired (or absent) slot is left for `replace` to overwrite — it + // guards on expiry and does the DEL+write atomically (single Lua exec + // in Redis, single lock in memory). We deliberately do NOT drop the + // stale entry here first: a standalone DEL is a separate, non-atomic op + // that can interleave between a sibling pod's `get` and its `replace`, + // clobbering a lease that pod just installed and orphaning it until + // server-side expiry. Reading a stale entry in the gap is harmless — + // every mutate/read path (`tryReserve`, `getCreditBalance`) already + // guards on expiry. The server treats an expired lease as released and + // refunds the full grant back to the company balance. + + const key = leaseKey(companyId, creditTypeId); + const inflight = this.inflightAcquire.get(key); + if (inflight) return inflight; + + const promise = this.acquire(companyId, creditTypeId).finally(() => { + this.inflightAcquire.delete(key); + }); + this.inflightAcquire.set(key, promise); + return promise; + } + + private async acquire(companyId: string, creditTypeId: string): Promise { + const resolved = this.resolveConfig(creditTypeId); + const body: api.AcquireCreditLeaseRequestBody = { + companyId, + creditTypeId, + requestedAmount: resolved.leaseSize, + expiresAt: new Date(Date.now() + resolved.leaseDuration), + }; + + try { + const response = await this.creditsClient.acquireCreditLease(body); + const data = response.data; + const wrote = await this.leaseStore.replace({ + leaseId: data.id, + companyId: data.companyId, + creditTypeId: data.creditTypeId, + grantedAmount: data.grantedAmount, + expiresAt: data.expiresAt, + }); + if (!wrote) { + // Another instance (sharing the backend) installed a live lease + // for this slot first — `replace` kept theirs to preserve its + // already-debited balance. The lease we just minted is therefore + // redundant. Release it so it isn't an orphaned hold against the + // company's balance until its server-side expiry. Fire-and-forget; + // a failed release just falls back to lease expiry. + this.logger.debug( + `Lost acquire race for ${companyId}/${creditTypeId}; releasing redundant lease ${data.id}`, + ); + void this.creditsClient + .releaseCreditLease(data.id, {}) + .catch((err) => + this.logger.warn(`Failed to release redundant credit lease ${data.id}: ${err}`), + ); + } else { + this.logger.debug( + `Acquired credit lease ${data.id} for ${companyId}/${creditTypeId} (granted=${data.grantedAmount}, expires=${data.expiresAt.toISOString()})`, + ); + } + return await this.leaseStore.get(companyId, creditTypeId); + } catch (err) { + this.logger.error(`Failed to acquire credit lease for ${companyId}/${creditTypeId}: ${err}`); + return undefined; + } + } + + /** + * Single-flight: kick off a background extend when one is warranted. + * An extend is triggered if EITHER: + * - the local remaining is below the low-water-mark ratio (steady-state + * refresh), or + * - the caller passes `requiredCredits` and the local remaining is below + * that figure (a single check just failed a reserve for that many + * credits — extend opportunistically instead of waiting for the next + * sub-watermark check). + * Returns the in-flight promise so callers can await it or fire-and-forget. + */ + async maybeExtendInBackground( + companyId: string, + creditTypeId: string, + requiredCredits?: number, + ): Promise { + const entry = await this.leaseStore.get(companyId, creditTypeId); + if (!entry) return undefined; + const resolved = this.resolveConfig(creditTypeId); + const ratio = entry.localRemainingCredits / Math.max(entry.grantedAmount, 1); + const belowWatermark = ratio <= resolved.lowWaterMark; + const belowRequired = requiredCredits !== undefined && entry.localRemainingCredits < requiredCredits; + if (!belowWatermark && !belowRequired) return entry; + + const key = leaseKey(companyId, creditTypeId); + const inflight = this.inflightExtend.get(key); + if (inflight) return inflight; + + const promise = this.extend(entry, resolved).finally(() => { + this.inflightExtend.delete(key); + }); + this.inflightExtend.set(key, promise); + return promise; + } + + private async extend(entry: LeaseEntry, resolved: ResolvedLeaseConfig): Promise { + const body: api.ExtendCreditLeaseRequestBody = { + additionalAmount: resolved.leaseSize, + expiresAt: new Date(Date.now() + resolved.leaseDuration), + }; + try { + const response = await this.creditsClient.extendCreditLease(entry.leaseId, body); + const data = response.data; + // The server returns the new totals; mirror them locally. + const delta = data.grantedAmount - entry.grantedAmount; + if (delta > 0) { + await this.leaseStore.extend(entry.companyId, entry.creditTypeId, delta, data.expiresAt); + } + this.logger.debug( + `Extended credit lease ${entry.leaseId} by ${delta} (now ${data.grantedAmount}, expires ${data.expiresAt.toISOString()})`, + ); + return await this.leaseStore.get(entry.companyId, entry.creditTypeId); + } catch (err) { + this.logger.warn(`Failed to extend credit lease ${entry.leaseId}: ${err}`); + return undefined; + } + } + +} diff --git a/src/credits/lease-store.ts b/src/credits/lease-store.ts new file mode 100644 index 00000000..fada3d94 --- /dev/null +++ b/src/credits/lease-store.ts @@ -0,0 +1,186 @@ +/** + * In-memory lease store, keyed by `${companyId}:${creditTypeId}`. + * + * Holds the lease ID, original granted amount, expiry, and the SDK's local + * view of `localRemainingCredits` — the portion of the lease not yet reserved + * by an outstanding reservation. Per-key serialization (Promise chain) keeps + * reserve / refund / replace atomic without external locking. + * + * The store doesn't talk to the API directly; `CreditLeaseManager` drives + * acquire/extend/release through the wire client and uses these methods to + * mirror remote state locally. + * + * For cross-pod deployments, swap this for `RedisLeaseStore` — both + * implement `ILeaseStore`. + */ + +export interface LeaseEntry { + leaseId: string; + companyId: string; + creditTypeId: string; + grantedAmount: number; + localRemainingCredits: number; + expiresAt: Date; +} + +/** Backing-store contract shared by `LeaseStore` (in-memory) and `RedisLeaseStore` (shared). */ +export interface ILeaseStore { + get(companyId: string, creditTypeId: string): Promise | LeaseEntry | undefined; + /** + * Install a fresh lease for the slot, but only if no live lease already + * occupies it. Implementations must be safe against concurrent writers + * sharing the backend: if a live (unexpired) lease is already present — + * even one with a different `leaseId`, e.g. acquired by a sibling pod that + * raced this one — leave it untouched so its already-debited + * `localRemainingCredits` wins. Returns `true` if it wrote a fresh row, + * `false` if it kept an existing live lease (the caller is then holding a + * redundant lease it should release). + */ + replace(entry: Omit): Promise; + extend( + companyId: string, + creditTypeId: string, + additionalGranted: number, + newExpiresAt?: Date, + ): Promise; + drop(companyId: string, creditTypeId: string): Promise; + tryReserve(companyId: string, creditTypeId: string, credits: number): Promise; + refund(companyId: string, creditTypeId: string, credits: number): Promise; +} + +export function leaseKey(companyId: string, creditTypeId: string): string { + return `${companyId}:${creditTypeId}`; +} + +export class LeaseStore implements ILeaseStore { + private leases = new Map(); + /** + * Per-key serializer chain. Each `withLock` appends to the tail so + * reserve / refund / replace see consistent intermediate state. + */ + private locks = new Map>(); + + /** Snapshot of the current entry, or undefined if no active lease. */ + get(companyId: string, creditTypeId: string): LeaseEntry | undefined { + const entry = this.leases.get(leaseKey(companyId, creditTypeId)); + return entry ? { ...entry } : undefined; + } + + /** + * Replace (or insert) the lease for the given (company, creditType). + * Resets `localRemainingCredits` to `grantedAmount` for a new lease. + * If a live lease already occupies the slot — regardless of `leaseId` — it + * is left alone so any already-debited `localRemainingCredits` wins. + * Returns `true` if a fresh row was written, `false` if an existing live + * lease was kept. + */ + async replace(entry: Omit): Promise { + const key = leaseKey(entry.companyId, entry.creditTypeId); + return this.withLock(key, () => { + const existing = this.leases.get(key); + if (existing && existing.expiresAt.getTime() > Date.now()) { + // A live lease already holds this slot — preserve its + // already-debited localRemaining instead of clobbering it. + return false; + } + this.leases.set(key, { + ...entry, + localRemainingCredits: entry.grantedAmount, + }); + return true; + }); + } + + /** + * Add granted credits to an existing lease (after a remote extend). If + * no lease exists for the key, this is a no-op (caller should have + * replaced). + */ + async extend( + companyId: string, + creditTypeId: string, + additionalGranted: number, + newExpiresAt?: Date, + ): Promise { + const key = leaseKey(companyId, creditTypeId); + return this.withLock(key, () => { + const entry = this.leases.get(key); + if (!entry) return; + entry.grantedAmount += additionalGranted; + entry.localRemainingCredits += additionalGranted; + if (newExpiresAt) entry.expiresAt = newExpiresAt; + }); + } + + /** Drop the lease entry (after a remote release or lazy expiry). */ + async drop(companyId: string, creditTypeId: string): Promise { + const key = leaseKey(companyId, creditTypeId); + return this.withLock(key, () => { + this.leases.delete(key); + }); + } + + /** + * Attempt to reserve `credits` from the local remaining balance. + * Returns true on success (debits the local balance), false if there + * isn't enough remaining. + */ + async tryReserve( + companyId: string, + creditTypeId: string, + credits: number, + ): Promise { + if (credits <= 0) return true; + const key = leaseKey(companyId, creditTypeId); + return this.withLock(key, () => { + const entry = this.leases.get(key); + if (!entry) return false; + // Never reserve against an expired lease: the server treats it as + // released and refunds the grant to the company balance, so the + // local `localRemainingCredits` is stale. Mirrors the expiry guard + // in the Redis store and the lazy expiry in `CreditLeaseManager`. + if (entry.expiresAt.getTime() <= Date.now()) return false; + if (entry.localRemainingCredits < credits) return false; + entry.localRemainingCredits -= credits; + return true; + }); + } + + /** Refund credits back to the local balance (capped at grantedAmount). */ + async refund( + companyId: string, + creditTypeId: string, + credits: number, + ): Promise { + if (credits <= 0) return; + const key = leaseKey(companyId, creditTypeId); + return this.withLock(key, () => { + const entry = this.leases.get(key); + if (!entry) return; + entry.localRemainingCredits = Math.min( + entry.localRemainingCredits + credits, + entry.grantedAmount, + ); + }); + } + + private async withLock(key: string, fn: () => T | Promise): Promise { + const prev = this.locks.get(key) ?? Promise.resolve(); + let resolve!: () => void; + const next = new Promise((r) => (resolve = r)); + const chain = prev.then(() => next); + this.locks.set(key, chain); + await prev; + try { + return await fn(); + } finally { + resolve(); + // Best-effort cleanup so the map doesn't grow without bound. + // Only delete if we're still the tail — otherwise a later caller + // has appended and needs the chain to stay live. + if (this.locks.get(key) === chain) { + this.locks.delete(key); + } + } + } +} diff --git a/src/credits/redis-lease-store.ts b/src/credits/redis-lease-store.ts new file mode 100644 index 00000000..e35bc64b --- /dev/null +++ b/src/credits/redis-lease-store.ts @@ -0,0 +1,223 @@ +import type { RedisClient } from "../cache/redis"; + +import { type ILeaseStore, type LeaseEntry, leaseKey } from "./lease-store"; +import { DEFAULT_LEASE_DURATION_MS } from "./types"; + +const DEFAULT_KEY_PREFIX = "schematic:"; +const LEASE_KEY_NAMESPACE = "credit-lease:"; +// How long after the declared expiry to keep the Redis row around before +// auto-eviction. Gives the sweeper a window to refund expired reservations +// before the underlying lease state disappears. +const LEASE_TTL_GRACE_MS = 60_000; + +// Every Lua script below touches exactly ONE key (the lease hash). That keeps +// them safe under Redis Cluster: a multi-key script whose keys hash to +// different slots raises CROSSSLOT, so cross-key atomicity is intentionally +// avoided here. The lease hash is the only state that must mutate atomically +// (check-and-decrement on reserve, clamped refund); index/bookkeeping that +// would span keys is done with ordinary single-key commands instead. +// +// Expiry is decided against the *Redis server's* clock (`redis.call('TIME')`), +// not the calling pod's `Date.now()`. With many pods sharing one lease, reading +// the wall clock locally would let clock skew disagree on whether a lease is +// still live: a fast pod denies reservations against a perfectly good lease +// (churn), a lagging pod reserves against a lease the server already swept and +// refunded (bounded over-spend). Sourcing "now" from Redis gives every pod a +// single authoritative clock. `redis.replicate_commands()` is called first so +// the non-deterministic TIME read is allowed alongside the script's writes on +// Redis 5/6 (it's a harmless no-op on 7+, which always uses effects +// replication). The `LEASE_NOW_MS` snippet converts TIME's `[sec, micros]` to +// integer milliseconds to match the `expiresAt` we store. +const LEASE_NOW_MS = ` +redis.replicate_commands() +local t = redis.call('TIME') +local now = (tonumber(t[1]) * 1000) + math.floor(tonumber(t[2]) / 1000) +`; + +/** + * Atomic `replace`. Writes the lease hash only when the slot is empty or the + * existing lease has expired. Returns "1" on write, "0" if a *live* lease + * already occupies the slot — even one with a different `leaseId`, e.g. + * installed by a sibling instance that raced this acquire. Keeping the existing + * live lease preserves its already-debited `localRemainingCredits`; the caller + * is left holding a redundant lease it should release. `now` is the Redis + * server clock (see `LEASE_NOW_MS`). + */ +const REPLACE_SCRIPT = + LEASE_NOW_MS + + ` +local existing_id = redis.call('HGET', KEYS[1], 'leaseId') +local existing_expiry = tonumber(redis.call('HGET', KEYS[1], 'expiresAt') or '0') +local new_id = ARGV[1] +local new_granted = ARGV[2] +local new_expiry = tonumber(ARGV[3]) +local grace = tonumber(ARGV[4]) + +if existing_id and existing_expiry > now then + return 0 +end + +redis.call('DEL', KEYS[1]) +redis.call('HSET', KEYS[1], + 'leaseId', new_id, + 'companyId', ARGV[5], + 'creditTypeId', ARGV[6], + 'grantedAmount', new_granted, + 'localRemainingCredits', new_granted, + 'expiresAt', ARGV[3]) +redis.call('PEXPIREAT', KEYS[1], new_expiry + grace) +return 1 +`; + +/** + * Atomic check-and-decrement on `localRemainingCredits`. Returns "1" on + * success, "0" if there's no lease, the lease has expired, or there's + * insufficient remaining. The expiry guard compares against the Redis server + * clock (`now`, see `LEASE_NOW_MS`) so a reserve against an expired-but-not-yet- + * evicted row during the TTL grace window is rejected — the server treats an + * expired lease as released, so its balance is stale. + */ +const TRY_RESERVE_SCRIPT = + LEASE_NOW_MS + + ` +local raw = redis.call('HGET', KEYS[1], 'localRemainingCredits') +if not raw then return 0 end +local expiry = tonumber(redis.call('HGET', KEYS[1], 'expiresAt') or '0') +if expiry <= now then return 0 end +local remaining = tonumber(raw) +local requested = tonumber(ARGV[1]) +if remaining < requested then return 0 end +redis.call('HSET', KEYS[1], 'localRemainingCredits', tostring(remaining - requested)) +return 1 +`; + +/** Refund credits, clamped at `grantedAmount`. */ +const REFUND_SCRIPT = ` +local raw_remaining = redis.call('HGET', KEYS[1], 'localRemainingCredits') +if not raw_remaining then return 0 end +local remaining = tonumber(raw_remaining) +local granted = tonumber(redis.call('HGET', KEYS[1], 'grantedAmount') or '0') +local refund = tonumber(ARGV[1]) +local new_balance = remaining + refund +if new_balance > granted then new_balance = granted end +redis.call('HSET', KEYS[1], 'localRemainingCredits', tostring(new_balance)) +return 1 +`; + +/** Bump grantedAmount + localRemainingCredits, update expiry. */ +const EXTEND_SCRIPT = ` +local raw_granted = redis.call('HGET', KEYS[1], 'grantedAmount') +if not raw_granted then return 0 end +local granted = tonumber(raw_granted) +local remaining = tonumber(redis.call('HGET', KEYS[1], 'localRemainingCredits') or '0') +local add = tonumber(ARGV[1]) +local new_expiry = tonumber(ARGV[2]) +local grace = tonumber(ARGV[3]) +redis.call('HSET', KEYS[1], + 'grantedAmount', tostring(granted + add), + 'localRemainingCredits', tostring(remaining + add), + 'expiresAt', ARGV[2]) +redis.call('PEXPIREAT', KEYS[1], new_expiry + grace) +return 1 +`; + +/** + * Redis-backed lease store. One hash per `(companyId, creditTypeId)` slot, with + * atomic mutations via single-key Lua scripts so it stays correct on both + * standalone and clustered Redis. + */ +export class RedisLeaseStore implements ILeaseStore { + private readonly client: RedisClient; + private readonly keyPrefix: string; + private readonly defaultLeaseDurationMs: number; + + constructor(opts: { + client: RedisClient; + keyPrefix?: string; + /** + * Defensive fallback used by `extend()` when the caller doesn't supply + * `newExpiresAt`. The lease manager always passes one today, so this + * only matters for direct callers. Default `DEFAULT_LEASE_DURATION_MS`. + */ + defaultLeaseDurationMs?: number; + }) { + this.client = opts.client; + this.keyPrefix = opts.keyPrefix ?? DEFAULT_KEY_PREFIX; + this.defaultLeaseDurationMs = opts.defaultLeaseDurationMs ?? DEFAULT_LEASE_DURATION_MS; + } + + /** Public so the reservation store can target the same lease hash for refunds. */ + hashKey(companyId: string, creditTypeId: string): string { + return `${this.keyPrefix}${LEASE_KEY_NAMESPACE}${leaseKey(companyId, creditTypeId)}`; + } + + async get(companyId: string, creditTypeId: string): Promise { + const raw = await this.client.hGetAll(this.hashKey(companyId, creditTypeId)); + if (!raw || !raw.leaseId) return undefined; + return decodeEntry(raw); + } + + async replace(entry: Omit): Promise { + const result = await this.client.eval(REPLACE_SCRIPT, { + keys: [this.hashKey(entry.companyId, entry.creditTypeId)], + // No client clock here — the script reads `now` from the Redis + // server via TIME (see LEASE_NOW_MS) so all pods agree on expiry. + arguments: [ + entry.leaseId, + String(entry.grantedAmount), + String(entry.expiresAt.getTime()), + String(LEASE_TTL_GRACE_MS), + entry.companyId, + entry.creditTypeId, + ], + }); + return Number(result) === 1; + } + + async extend( + companyId: string, + creditTypeId: string, + additionalGranted: number, + newExpiresAt?: Date, + ): Promise { + const expiry = (newExpiresAt ?? new Date(Date.now() + this.defaultLeaseDurationMs)).getTime(); + await this.client.eval(EXTEND_SCRIPT, { + keys: [this.hashKey(companyId, creditTypeId)], + arguments: [String(additionalGranted), String(expiry), String(LEASE_TTL_GRACE_MS)], + }); + } + + async drop(companyId: string, creditTypeId: string): Promise { + // A plain single-key delete — no secondary index to keep in sync. + await this.client.del(this.hashKey(companyId, creditTypeId)); + } + + async tryReserve(companyId: string, creditTypeId: string, credits: number): Promise { + if (credits <= 0) return true; + const result = await this.client.eval(TRY_RESERVE_SCRIPT, { + keys: [this.hashKey(companyId, creditTypeId)], + // Only the requested amount — `now` comes from the Redis server clock. + arguments: [String(credits)], + }); + return Number(result) === 1; + } + + async refund(companyId: string, creditTypeId: string, credits: number): Promise { + if (credits <= 0) return; + await this.client.eval(REFUND_SCRIPT, { + keys: [this.hashKey(companyId, creditTypeId)], + arguments: [String(credits)], + }); + } +} + +function decodeEntry(raw: Record): LeaseEntry { + return { + leaseId: raw.leaseId, + companyId: raw.companyId, + creditTypeId: raw.creditTypeId, + grantedAmount: Number(raw.grantedAmount ?? "0"), + localRemainingCredits: Number(raw.localRemainingCredits ?? "0"), + expiresAt: new Date(Number(raw.expiresAt ?? "0")), + }; +} diff --git a/src/credits/redis-reservation-store.ts b/src/credits/redis-reservation-store.ts new file mode 100644 index 00000000..bb196161 --- /dev/null +++ b/src/credits/redis-reservation-store.ts @@ -0,0 +1,277 @@ +import type { RedisClient } from "../cache/redis"; + +import type { ILeaseStore } from "./lease-store"; +import type { IReservationStore } from "./reservation-store"; +import type { Reservation } from "./types"; + +const DEFAULT_KEY_PREFIX = "schematic:"; +const RES_KEY_NAMESPACE = "credit-reservation:"; +// Sorted set scoring open reservations by `expiresAt` so the sweeper can pop +// expired entries in O(log n). Members are NOT bare reservation ids: they +// encode `companyId|creditTypeId|id` (see `encodeMember`) so the sweeper can +// reconcile the secondary indexes — this set and the per-tenant `byCredit` +// hash — even when the reservation hash has already TTL-evicted and the CLAIM +// in `consume` can no longer report which (company, credit) the hold belonged +// to. Without that, an evicted-before-swept reservation would orphan its +// `byCredit` field forever and permanently inflate `reservedCredits`. +const RES_INDEX_KEY = "credit-reservations:byExpiry"; +// Per-(company, credit) index of open reservations, stored as a single hash +// of `reservationId -> creditsReserved`. `reservedCredits` then reads the whole +// tenant's holds with ONE `HGETALL` (single key, Cluster-safe) and sums the +// values, instead of fanning out a per-reservation read. The hash also *is* the +// source of truth for the sum: a field exists iff its reservation is open and +// unrefunded (consume removes the field and refunds in the same call), so the +// sum stays exact without cross-checking the reservation hashes. +const RES_BYCREDIT_NAMESPACE = "credit-reservations:byCredit:"; +// Buffer past `expiresAt` before Redis auto-evicts the row, so the sweeper +// has a window to refund. +const RES_TTL_GRACE_MS = 30_000; + +/** + * Atomic claim: read the reservation hash and delete it in one step, returning + * its fields (or `nil` if it was already gone). Touches a single key, so it's + * safe under Redis Cluster. The atomic read-then-delete is what makes + * `consume` exactly-once: of two racing callers (a normal Track and a sweeper, + * say) only one gets the fields back and proceeds to refund — the other sees + * `nil`. The refund to the lease hash is a separate single-key op; a crash in + * the gap leaves the unspent slice held on the lease until the lease itself + * expires, never double-refunded. + */ +const CLAIM_SCRIPT = ` +local raw = redis.call('HGETALL', KEYS[1]) +if #raw == 0 then return nil end +redis.call('DEL', KEYS[1]) +return raw +`; + +// `byExpiry` zset members encode the (company, credit, id) tuple so the sweeper +// can clean the per-tenant `byCredit` hash even after the reservation hash has +// TTL-evicted (at which point CLAIM returns nil and can't report company/credit). +// The delimiter `|` is absent from Schematic ids and the UUID reservation id. +function encodeMember(r: { companyId: string; creditTypeId: string; id: string }): string { + return `${r.companyId}|${r.creditTypeId}|${r.id}`; +} + +// Returns undefined for a legacy bare-id member (written by a pre-upgrade SDK +// still in the set during a rolling deploy); the sweeper falls back to treating +// the whole member as the id in that case. +function decodeMember(member: string): { companyId: string; creditTypeId: string; id: string } | undefined { + const parts = member.split("|"); + if (parts.length !== 3) return undefined; + return { companyId: parts[0], creditTypeId: parts[1], id: parts[2] }; +} + +/** + * Redis-backed reservation table. Each reservation is a hash, indexed by + * `expiresAt` in a sorted set so the sweeper can pop expired entries in + * O(log n), and by `(company, credit)` so the balance display can sum a + * tenant's open holds. All mutations use single-key operations (or single-key + * Lua), so the store is correct on standalone and clustered Redis alike — the + * unspent-slice refund is delegated to the lease store rather than reaching + * across to the lease hash inside a multi-key script. + */ +export class RedisReservationStore implements IReservationStore { + private readonly client: RedisClient; + private readonly leaseStore: ILeaseStore; + private readonly sweepIntervalMs: number; + private readonly keyPrefix: string; + private sweepInterval: NodeJS.Timeout | null = null; + private stopped = false; + + constructor(opts: { + client: RedisClient; + leaseStore: ILeaseStore; + sweepIntervalMs?: number; + keyPrefix?: string; + }) { + this.client = opts.client; + this.leaseStore = opts.leaseStore; + this.sweepIntervalMs = opts.sweepIntervalMs ?? 1000; + this.keyPrefix = opts.keyPrefix ?? DEFAULT_KEY_PREFIX; + } + + private hashKey(id: string): string { + return `${this.keyPrefix}${RES_KEY_NAMESPACE}${id}`; + } + + private indexKey(): string { + return `${this.keyPrefix}${RES_INDEX_KEY}`; + } + + private byCreditKey(companyId: string, creditTypeId: string): string { + return `${this.keyPrefix}${RES_BYCREDIT_NAMESPACE}${companyId}:${creditTypeId}`; + } + + async add(reservation: Reservation): Promise { + const expiresMs = reservation.expiresAt.getTime(); + const hashKey = this.hashKey(reservation.id); + // Write the hash (+ TTL) first so the reservation exists before anything + // references it, then index it for the sweeper and the per-tenant sum. + // These are independent single-key ops rather than one multi-key script: + // a partial failure at worst leaves an un-indexed reservation that the + // TTL reaps (its slice reclaimed when the lease expires), never a + // double-spend. + await this.client.hSet(hashKey, { + id: reservation.id, + leaseId: reservation.leaseId, + companyId: reservation.companyId, + creditTypeId: reservation.creditTypeId, + eventSubtype: reservation.eventSubtype, + quantityReserved: String(reservation.quantityReserved), + creditsReserved: String(reservation.creditsReserved), + consumptionRate: String(reservation.consumptionRate), + expiresAt: String(expiresMs), + evalCtx: JSON.stringify(reservation.evalCtx), + }); + await this.client.pExpireAt(hashKey, expiresMs + RES_TTL_GRACE_MS); + await this.client.zAdd(this.indexKey(), { score: expiresMs, value: encodeMember(reservation) }); + // Record the credits in the per-tenant hash keyed by reservation id, so + // `reservedCredits` can sum the tenant's open holds in one HGETALL. + await this.client.hSet( + this.byCreditKey(reservation.companyId, reservation.creditTypeId), + reservation.id, + String(reservation.creditsReserved), + ); + } + + async get(id: string): Promise { + const raw = await this.client.hGetAll(this.hashKey(id)); + if (!raw || !raw.id) return undefined; + return decodeReservation(raw); + } + + /** + * Sum open reservations for a (company, credit) with a single `HGETALL` on + * the per-tenant index hash — one round trip, one key (Cluster-safe), no + * per-reservation fan-out. The hash values are the authoritative + * `creditsReserved` figures; a field is present iff its reservation is open + * and unrefunded (consume removes the field and refunds in the same call), + * so summing them is exact. Display-path call, not a hot path. + */ + async reservedCredits(companyId: string, creditTypeId: string): Promise { + const byCredit = await this.client.hGetAll(this.byCreditKey(companyId, creditTypeId)).catch(() => ({})); + let total = 0; + for (const value of Object.values(byCredit)) { + total += Number(value) || 0; + } + return total; + } + + async consume(id: string, creditsConsumed: number): Promise { + // Atomically claim (read + delete) the reservation hash. Only one caller + // wins; a duplicate/racing consume gets nil and returns null. + const claimed = await this.client.eval(CLAIM_SCRIPT, { + keys: [this.hashKey(id)], + arguments: [], + }); + const raw = decodeRawArray(claimed); + if (!raw || !raw.id) return null; + + const companyId = raw.companyId; + const creditTypeId = raw.creditTypeId; + const reserved = Number(raw.creditsReserved) || 0; + + // Index cleanup — single-key ops. Drop the credits from the per-tenant + // hash BEFORE the refund below so the lease (localRemaining + this hash) + // never transiently double-counts the slice: while it sits on the hash + // it's "reserved", and the refund moves it back to localRemaining. + await this.client.zRem(this.indexKey(), encodeMember({ companyId, creditTypeId, id })).catch(() => {}); + await this.client.hDel(this.byCreditKey(companyId, creditTypeId), id).catch(() => {}); + + let consumed = creditsConsumed; + if (consumed < 0) consumed = 0; + if (consumed > reserved) consumed = reserved; + const refund = reserved - consumed; + if (refund > 0) { + // Delegate the clamped refund to the lease store, which owns the + // lease hash. Keeps the cross-key write out of a single Lua script. + await this.leaseStore.refund(companyId, creditTypeId, refund); + } + return consumed; + } + + startSweep(): void { + if (this.sweepInterval || this.stopped) return; + this.sweepInterval = setInterval(() => { + this.sweepExpired().catch(() => { + // Swallow so the timer keeps running. + }); + }, this.sweepIntervalMs); + if (this.sweepInterval.unref) this.sweepInterval.unref(); + } + + async sweepExpired(now: Date = new Date()): Promise { + const cutoff = now.getTime(); + // ZRANGEBYSCORE returns members (encoded `company|credit|id`) whose score + // (expiresAt) is <= cutoff. + const expired = await this.client.zRangeByScore(this.indexKey(), 0, cutoff); + let swept = 0; + for (const member of expired) { + const decoded = decodeMember(member); + // Fall back to treating the whole member as the id for legacy + // bare-id members written before this SDK upgrade. + const id = decoded?.id ?? member; + const refunded = await this.consume(id, 0); + // Always drop the member we read. On the success path `consume` + // already removed the current-format member (so this is an idempotent + // no-op); it also covers a legacy bare-id member `consume` couldn't + // match, and the hash-evicted path below. + await this.client.zRem(this.indexKey(), member).catch(() => {}); + if (refunded !== null) { + swept++; + continue; + } + // `consume` found no reservation hash. Either a racing track already + // consumed it (and reconciled the `byCredit` field — the hDel below is + // then a no-op), or the hash TTL-evicted before the sweeper reached it, + // orphaning the `byCredit` field. Reconcile it so `reservedCredits` + // can't keep summing an evicted hold. We do NOT refund the unspent + // slice here: without the hash, CLAIM can't arbitrate exactly-once + // across racing sweepers, so the slice is reclaimed when the lease + // itself expires server-side instead. + if (decoded) { + await this.client + .hDel(this.byCreditKey(decoded.companyId, decoded.creditTypeId), decoded.id) + .catch(() => {}); + } + } + return swept; + } + + stop(): void { + this.stopped = true; + if (this.sweepInterval) { + clearInterval(this.sweepInterval); + this.sweepInterval = null; + } + } + + async size(): Promise { + return this.client.zCard(this.indexKey()).catch(() => 0); + } +} + +/** Decode a flat `[field, value, field, value, ...]` HGETALL array (as CLAIM_SCRIPT returns). */ +function decodeRawArray(raw: unknown): Record | undefined { + if (!Array.isArray(raw) || raw.length === 0) return undefined; + const out: Record = {}; + for (let i = 0; i + 1 < raw.length; i += 2) { + out[String(raw[i])] = String(raw[i + 1]); + } + return out; +} + +function decodeReservation(raw: Record): Reservation { + return { + id: raw.id, + leaseId: raw.leaseId, + companyId: raw.companyId, + creditTypeId: raw.creditTypeId, + eventSubtype: raw.eventSubtype, + quantityReserved: Number(raw.quantityReserved), + creditsReserved: Number(raw.creditsReserved), + consumptionRate: Number(raw.consumptionRate), + expiresAt: new Date(Number(raw.expiresAt)), + evalCtx: raw.evalCtx ? JSON.parse(raw.evalCtx) : {}, + }; +} diff --git a/src/credits/reservation-store.ts b/src/credits/reservation-store.ts new file mode 100644 index 00000000..95bfe0b1 --- /dev/null +++ b/src/credits/reservation-store.ts @@ -0,0 +1,137 @@ +import type { Reservation } from "./types"; +import type { ILeaseStore } from "./lease-store"; + +/** Backing-store contract for the reservation table. */ +export interface IReservationStore { + add(reservation: Reservation): Promise | void; + get(id: string): Promise | Reservation | undefined; + consume(id: string, creditsConsumed: number): Promise; + /** + * Sum of `creditsReserved` across open reservations for a (company, credit). + * These credits are carved out of the lease's `localRemainingCredits` but + * aren't spent yet, so a balance display can add them back to ignore + * in-flight holds. + */ + reservedCredits(companyId: string, creditTypeId: string): Promise | number; + startSweep(): void; + sweepExpired(now?: Date): Promise; + stop(): void; + size(): Promise | number; +} + +/** + * In-memory reservation table, paired with a sweep loop that returns expired + * reservations to their underlying lease. + * + * For cross-pod deployments, swap this for `RedisReservationStore` — both + * implement `IReservationStore`. + */ +export class ReservationStore implements IReservationStore { + private reservations = new Map(); + private sweepInterval: NodeJS.Timeout | null = null; + private stopped = false; + + constructor( + private readonly leaseStore: ILeaseStore, + private readonly sweepIntervalMs: number = 1000, + ) {} + + /** Register a new reservation. Idempotent on `id`. */ + add(reservation: Reservation): void { + this.reservations.set(reservation.id, reservation); + } + + /** Look up a reservation by ID. */ + get(id: string): Reservation | undefined { + return this.reservations.get(id); + } + + /** + * Sum open reservations for a (company, credit). Counts every reservation + * still in the table: its credits stay carved out of the lease's + * `localRemainingCredits` until `consume`/`sweepExpired` removes it and + * refunds the unspent remainder in the same step, so summing by presence + * keeps `localRemainingCredits + reservedCredits` exact. + */ + reservedCredits(companyId: string, creditTypeId: string): number { + let total = 0; + for (const reservation of this.reservations.values()) { + if ( + reservation.companyId === companyId && + reservation.creditTypeId === creditTypeId + ) { + total += reservation.creditsReserved; + } + } + return total; + } + + /** + * Consume a reservation — removes it from the table and refunds + * `creditsReserved - creditsConsumed` back to the underlying lease. + * Returns the credits actually consumed (clamped to `creditsReserved`) + * or `null` if the reservation is missing/already consumed. + */ + async consume(id: string, creditsConsumed: number): Promise { + const reservation = this.reservations.get(id); + if (!reservation) return null; + this.reservations.delete(id); + + const actual = Math.max(0, Math.min(creditsConsumed, reservation.creditsReserved)); + const refund = reservation.creditsReserved - actual; + if (refund > 0) { + await this.leaseStore.refund( + reservation.companyId, + reservation.creditTypeId, + refund, + ); + } + return actual; + } + + /** Start the background sweep loop. Safe to call repeatedly. */ + startSweep(): void { + if (this.sweepInterval || this.stopped) return; + this.sweepInterval = setInterval(() => { + this.sweepExpired().catch(() => { + // sweepExpired only throws on programmer error; swallow so the + // timer keeps running. + }); + }, this.sweepIntervalMs); + if (this.sweepInterval.unref) this.sweepInterval.unref(); + } + + /** + * Scan for expired reservations, remove them, and refund their credits to + * the lease. + */ + async sweepExpired(now: Date = new Date()): Promise { + let swept = 0; + for (const [id, reservation] of this.reservations) { + if (reservation.expiresAt.getTime() <= now.getTime()) { + this.reservations.delete(id); + await this.leaseStore.refund( + reservation.companyId, + reservation.creditTypeId, + reservation.creditsReserved, + ); + swept++; + } + } + return swept; + } + + /** Stop the sweep loop. */ + stop(): void { + this.stopped = true; + if (this.sweepInterval) { + clearInterval(this.sweepInterval); + this.sweepInterval = null; + } + } + + /** Test/debug: current reservation count. */ + size(): number { + return this.reservations.size; + } +} diff --git a/src/credits/track.ts b/src/credits/track.ts new file mode 100644 index 00000000..90a9a02f --- /dev/null +++ b/src/credits/track.ts @@ -0,0 +1,62 @@ +import * as api from "../api"; + +import type { IReservationStore } from "./reservation-store"; +import type { Reservation, TrackWithReservationOptions } from "./types"; + +/** Outcome of consuming a reservation and building its Track event. */ +export interface ReservationConsumeResult { + /** The Track event to emit. Always carries the `leaseId`. */ + track: api.EventBodyTrack; + /** + * `true` when the reservation was still live and `consume` debited/refunded + * the local lease in this call. `false` when the reservation was already + * gone from the store (swept after its TTL, or already consumed): the local + * lease balance was *not* touched here, and the returned `track` is a + * recovery event so the server still bills the actual usage. The caller is + * responsible for de-duping recovery emits (see `trackWithReservation`). + */ + settledLocally: boolean; +} + +/** + * Build the `EventBodyTrack` payload for a Track event that consumes a + * reservation, and consume the reservation against the lease (refunding the + * unused slice locally). + * + * The Track is built from the caller-held `reservation` handle, so it can be + * emitted even when the reservation has already been swept out of the store — + * the server is the source of truth for actual consumption. When the lease is + * still live server-side the `leaseId` routes the spend through its sub-ledger; + * when the server lease has itself expired/released, the server falls through + * to a direct grant decrement. Either way the usage is billed, so a hold that + * outlives its (client-side) reservation TTL is no longer dropped on the floor. + */ +export async function consumeReservationAndBuildEvent( + reservations: IReservationStore, + reservation: Reservation, + actualQuantity: number, + options?: TrackWithReservationOptions, +): Promise { + const credits = actualQuantity * reservation.consumptionRate; + const consumed = await reservations.consume(reservation.id, credits); + + const body: api.EventBodyTrack = { + event: reservation.eventSubtype, + quantity: actualQuantity, + // Routes the server-side credit consumption through the lease's + // sub-ledger instead of decrementing the grant again (which was + // already pre-debited at acquire/extend). Without this the grant + // double-debits and eventually starves redemptions mid-session. + leaseId: reservation.leaseId, + }; + if (reservation.evalCtx.company) { + body.company = reservation.evalCtx.company; + } + if (reservation.evalCtx.user) { + body.user = reservation.evalCtx.user; + } + if (options?.traits) { + body.traits = options.traits; + } + return { track: body, settledLocally: consumed !== null }; +} diff --git a/src/credits/types.ts b/src/credits/types.ts new file mode 100644 index 00000000..6c0a9f50 --- /dev/null +++ b/src/credits/types.ts @@ -0,0 +1,176 @@ +import type * as api from "../api"; +import type { RedisClient } from "../cache/redis"; + +/** + * Behavior when a lease cannot be acquired (e.g. API error, insufficient balance). + * - `fail-open`: `check()` returns `{ allowed: true }` so the caller proceeds without preflight gating + * - `fail-closed`: `check()` returns `{ allowed: false }` so the caller blocks the action + */ +export type OnAcquireFailure = "fail-open" | "fail-closed"; + +// Single source of truth for defaults consumed by `CreditLeaseManager` and +// `SchematicClient`. Re-exported so the values referenced in the doc strings +// below stay accurate even if a consumer overrides only a subset of fields. +export const DEFAULT_LEASE_DURATION_MS: number = 5 * 60 * 1000; +export const DEFAULT_RESERVATION_TTL_MS: number = 60 * 1000; +export const DEFAULT_LEASE_SIZE: number = 10_000; +export const DEFAULT_LOW_WATER_MARK: number = 0.25; +export const DEFAULT_SWEEP_INTERVAL_MS: number = 1000; +// How long `prewarm()` is willing to wait for a freshly-identified company to +// surface in the datastream cache before giving up. Long enough to cover the +// buffer-flush → server-ingest → datastream-push round-trip for a new +// company; short enough that a misconfigured caller doesn't hang. +export const DEFAULT_PREWARM_RESOLVE_TIMEOUT_MS: number = 5000; +export const DEFAULT_PREWARM_POLL_INTERVAL_MS: number = 100; + +/** + * Configuration block enabling client-side lease + reservation behavior on + * `client.check` / `client.trackWithReservation`. Omit to keep the SDK + * lease-unaware (`check` falls back to a plain flag check). + */ +export interface CreditLeaseConfig { + /** Default lease duration in milliseconds. Default `DEFAULT_LEASE_DURATION_MS` (5 minutes). */ + defaultLeaseDuration?: number; + /** Default reservation TTL in milliseconds. Default `DEFAULT_RESERVATION_TTL_MS` (60 seconds). */ + defaultReservationTTL?: number; + /** Default lease size (credit amount requested). Default `DEFAULT_LEASE_SIZE` (10000). */ + defaultLeaseSize?: number; + /** + * Fraction of remaining lease balance below which the SDK kicks off a + * background extend. Default `DEFAULT_LOW_WATER_MARK` (0.25). + */ + lowWaterMark?: number; + /** Sweep interval (ms) for expired reservations. Default `DEFAULT_SWEEP_INTERVAL_MS` (1000). */ + sweepIntervalMs?: number; + /** + * Max time `prewarm()` will wait for a freshly-identified company to + * surface in the datastream cache when only secondary keys are passed. + * Set to 0 to skip waiting entirely (prewarm bails immediately if the + * company isn't already cached). Default `DEFAULT_PREWARM_RESOLVE_TIMEOUT_MS` + * (5000ms). + */ + prewarmResolveTimeoutMs?: number; + /** + * Pre-connected Redis client for lease + reservation state. Optional: when + * omitted, the SDK reuses the DataStream cache's Redis client + * (`dataStream.redisClient`) if one is configured, so an existing Redis + * setup backs leases automatically. Set this only to point lease state at a + * *different* Redis than the DataStream cache. + * + * When neither this nor `dataStream.redisClient` is configured, the SDK + * falls back to per-process in-memory stores — single-pod only, since + * cross-pod gating is lost (a warning is logged). With a Redis client + * present, the Lua-driven `tryReserve` / `consume` paths give atomic + * cross-pod gating. + */ + redisClient?: RedisClient; + /** + * Optional Redis key prefix. Falls back to `dataStream.redisKeyPrefix` when + * unset, then to `schematic:`. + */ + redisKeyPrefix?: string; + /** Per-credit-type overrides (keyed by credit type ID). */ + overrides?: Record>>; +} + +/** Resolved config for a single credit type after applying defaults + overrides. */ +export interface ResolvedLeaseConfig { + leaseDuration: number; + reservationTTL: number; + leaseSize: number; + lowWaterMark: number; +} + +/** Handle returned by `client.check` when a reservation is issued. Pass to `client.trackWithReservation`. */ +export interface Reservation { + /** Opaque reservation ID. */ + id: string; + /** Underlying lease ID this reservation draws from. */ + leaseId: string; + /** Company that owns the lease. */ + companyId: string; + /** Credit type the reservation reserves against. */ + creditTypeId: string; + /** Event subtype to record on the Track event when consumed. */ + eventSubtype: string; + /** Quantity (in event units) the caller declared upfront. */ + quantityReserved: number; + /** Credits reserved = `quantityReserved * consumptionRate`. */ + creditsReserved: number; + /** Consumption rate at the time the reservation was issued. */ + consumptionRate: number; + /** When the reservation expires and gets swept back to the lease. */ + expiresAt: Date; + /** + * Evaluation context used to issue this reservation. Threaded into the + * Track event in `trackWithReservation` so the server can attribute usage + * to the same company/user. + */ + evalCtx: api.CheckFlagRequestBody; +} + +/** Options accepted by `client.check`. */ +export interface CheckOptions { + /** + * Per-event-subtype simulated quantity. Pass + * `{ "inference_tokens": maxTokens }` to gate the check on a + * `maxTokens × consumption_rate` slice of the lease balance. + */ + eventUsage?: Record; + /** + * Single integer quantity applied to whatever numeric condition is being + * evaluated. Less specific than `eventUsage`; use when the event subtype + * is unambiguous from the flag. + */ + usage?: number; + /** + * What to do when a lease cannot be acquired (API error, insufficient + * remote balance). Default: `fail-closed` (`allowed = false`, no + * reservation) — deny when the gate can't gate. Override to `fail-open` for + * trusted/known customers where letting traffic through is preferable to a + * denial. + */ + onAcquireFailure?: OnAcquireFailure; + /** Default value to return on error. */ + defaultValue?: boolean | (() => boolean); + /** Custom timeout for API calls within this check (ms). */ + timeoutMs?: number; +} + +/** Result of `client.check`. */ +export interface CheckResult { + /** Whether the caller is permitted to proceed. */ + allowed: boolean; + /** Boolean flag value (`allowed` mirrors this in non-lease paths). */ + value: boolean; + /** Reservation handle when a lease-bearing check passed. */ + reservation?: Reservation; + /** Human-readable reason (from the rules engine or the SDK). */ + reason: string; + /** + * Entitlement payload from the check. + * + * NOTE: `entitlement.creditRemaining` (and `creditTotal` / `creditUsed`) is + * NOT lease-aware and must not be used as a user-facing balance when credit + * leases are enabled. The WASM derives those fields from the company's + * server `creditBalances`, which a lease distorts: a plain `checkFlag` sees + * the balance net of the whole lease tranche (reads low/~0), and a + * lease-bearing `check()` sees the substituted *tranche-local* balance, not + * the company total. For any "credits remaining" display use + * `client.getCreditBalance(...).settled` (`B − spent`), which speaks the + * server's `remaining`/`reserved`/`settled` vocabulary. + */ + entitlement?: api.RulesengineFeatureEntitlement; + /** Flag key checked. */ + flagKey: string; + /** Flag ID if known. */ + flagId?: string; + /** Optional error string (populated when the SDK fell back to a default). */ + err?: string; +} + +/** Extras accepted by `trackWithReservation`. */ +export interface TrackWithReservationOptions { + /** Optional traits to attach to the emitted Track event. */ + traits?: Record; +} diff --git a/src/datastream/datastream-client.ts b/src/datastream/datastream-client.ts index 0a5dfb35..4ca3000a 100644 --- a/src/datastream/datastream-client.ts +++ b/src/datastream/datastream-client.ts @@ -2,6 +2,7 @@ import * as Schematic from '../api/types'; import type { DatastreamWSClient } from './websocket-client'; import { DataStreamResp, DataStreamReq, DataStreamError, EntityType, MessageType } from './types'; import { RulesEngineClient } from '../rules-engine'; +import type { CheckFlagOptions } from '../wrapper'; import { Logger } from '../logger'; import { LazyEmitter } from './emitter'; import { partialCompany, partialUser, deepCopyCompany as deepCopyCompanyFn } from './merge'; @@ -475,7 +476,8 @@ export class DataStreamClient extends LazyEmitter { */ public async checkFlag( evalCtx: { company?: Record; user?: Record }, - flagKey: string + flagKey: string, + options?: CheckFlagOptions ): Promise { // Get flag first - return error if not found const flag = await this.getFlag(flagKey); @@ -503,13 +505,13 @@ export class DataStreamClient extends LazyEmitter { if (this.replicatorMode) { // In replicator mode, if we don't have all cached data, evaluate with null values instead of fetching // The external replicator should have populated the cache with all necessary data - return this.evaluateFlag(flag, cachedCompany, cachedUser); + return this.evaluateFlag(flag, cachedCompany, cachedUser, options); } // Non-replicator mode: if we have all cached data we need, use it if ((!needsCompany || cachedCompany) && (!needsUser || cachedUser)) { this.logger.debug(`All required resources found in cache for flag ${flagKey} evaluation`); - return this.evaluateFlag(flag, cachedCompany, cachedUser); + return this.evaluateFlag(flag, cachedCompany, cachedUser, options); } // Check if we're connected to datastream for live fetching @@ -529,7 +531,30 @@ export class DataStreamClient extends LazyEmitter { const [company, user] = await Promise.all([companyPromise, userPromise]); // Evaluate against the rules engine - return this.evaluateFlag(flag, company, user); + return this.evaluateFlag(flag, company, user, options); + } + + /** + * Public accessor for the cached company snapshot. Used by credit-lease + * callers that need the company's `credit_balances` so they can substitute + * a lease-bound view before calling the rules engine. + */ + public async getCachedCompany( + keys: Record, + ): Promise { + return this.getCompanyFromCache(keys); + } + + /** Public accessor for the cached user snapshot. */ + public async getCachedUser( + keys: Record, + ): Promise { + return this.getUserFromCache(keys); + } + + /** Public accessor for the rules engine. Used by credit-lease eval. */ + public getRulesEngine(): RulesEngineClient { + return this.rulesEngine; } /** @@ -1354,7 +1379,8 @@ export class DataStreamClient extends LazyEmitter { private async evaluateFlag( flag: Schematic.RulesengineFlag, company: Schematic.RulesengineCompany | null, - user: Schematic.RulesengineUser | null + user: Schematic.RulesengineUser | null, + options?: CheckFlagOptions ): Promise { const defaultValue = flag.defaultValue ?? false; @@ -1363,7 +1389,7 @@ export class DataStreamClient extends LazyEmitter { if (this.rulesEngine.isInitialized()) { this.logger.debug(`Evaluating flag with rules engine: ${JSON.stringify({ flagId: flag.id, flagRules: flag.rules?.length || 0, companyId: company?.id, userId: user?.id })}`); - const result = await this.rulesEngine.checkFlag(flag, company, user); + const result = await this.rulesEngine.checkFlagWithOptions(flag, company, user, options); this.logger.debug(`Rules engine evaluation result: ${JSON.stringify(result)}`); return { diff --git a/src/index.ts b/src/index.ts index 4359e33b..f1fae6d5 100644 --- a/src/index.ts +++ b/src/index.ts @@ -12,6 +12,14 @@ export { ConsoleLogger, LogLevel, type Logger } from "./logger"; export { SchematicEnvironment } from "./environments"; export { SchematicError, SchematicTimeoutError } from "./errors"; export { RulesEngineClient } from "./rules-engine"; +export type { + CheckOptions, + CheckResult, + CreditLeaseConfig, + OnAcquireFailure, + Reservation, + TrackWithReservationOptions, +} from "./credits"; export { verifyWebhookSignature, verifySignature, diff --git a/src/rules-engine.ts b/src/rules-engine.ts index a5337e46..62407ebc 100644 --- a/src/rules-engine.ts +++ b/src/rules-engine.ts @@ -1,4 +1,5 @@ import * as Schematic from './api/types'; +import type { CheckFlagOptions } from './wrapper'; /** Entitlement details returned by the WASM rules engine */ export interface WasmFeatureEntitlement { @@ -70,6 +71,15 @@ export class RulesEngineClient { flag: object, company?: object | null, user?: object | null + ): Promise { + return this.checkFlagWithOptions(flag, company, user); + } + + async checkFlagWithOptions( + flag: object, + company?: object | null, + user?: object | null, + options?: CheckFlagOptions | null ): Promise { this.ensureInitialized(); @@ -81,11 +91,13 @@ export class RulesEngineClient { const flagJson = JSON.stringify(flag, stripNulls); const companyJson = company ? JSON.stringify(company, stripNulls) : undefined; const userJson = user ? JSON.stringify(user, stripNulls) : undefined; + const optionsJson = options ? JSON.stringify(serializeCheckFlagOptions(options), stripNulls) : undefined; - const resultJson = this.wasmInstance!.checkFlag( + const resultJson = this.wasmInstance!.checkFlagWithOptions( flagJson, companyJson, - userJson + userJson, + optionsJson ); return JSON.parse(resultJson); @@ -117,5 +129,20 @@ export class RulesEngineClient { } } +// Serialize CheckFlagOptions to the snake_case envelope the WASM expects. +function serializeCheckFlagOptions(options: CheckFlagOptions): Record { + const envelope: Record = {}; + if (options.creditCost && Object.keys(options.creditCost).length > 0) { + envelope.credit_cost = options.creditCost; + } + if (options.usage !== undefined) { + envelope.usage = options.usage; + } + if (options.eventUsage && Object.keys(options.eventUsage).length > 0) { + envelope.event_usage = options.eventUsage; + } + return envelope; +} + // Export for backward compatibility export { RulesEngineClient as default }; \ No newline at end of file diff --git a/src/wrapper.ts b/src/wrapper.ts index f15975b9..33c62f31 100644 --- a/src/wrapper.ts +++ b/src/wrapper.ts @@ -9,6 +9,32 @@ import { offlineFetcher, provideFetcher } from "./core/fetcher/custom"; import { RUNTIME } from "./core/runtime"; import { DataStreamClient, type DataStreamClientOptions } from "./datastream"; import type { RedisClient } from "./cache/redis"; +import { + CreditLeaseManager, + DEFAULT_PREWARM_POLL_INTERVAL_MS, + DEFAULT_PREWARM_RESOLVE_TIMEOUT_MS, + DEFAULT_SWEEP_INTERVAL_MS, + LeaseStore, + RedisLeaseStore, + RedisReservationStore, + ReservationStore, + type ILeaseStore, + type IReservationStore, + type CheckOptions, + type CheckResult, + type CreditLeaseConfig, + type Reservation, + type TrackWithReservationOptions, +} from "./credits"; +import { checkWithLease } from "./credits/check"; +import { consumeReservationAndBuildEvent } from "./credits/track"; + +// Idempotency-key namespace for the Track event a reservation settles into. +// Deterministic per reservation, so a recovery emit (work outlived the local +// reservation TTL) and an accidental double `trackWithReservation` collapse to +// one event server-side: the events pipeline dedupes by (account, env, +// event_type, key) for 24h before any credit consumption runs. +const RESERVATION_TRACK_IDEMPOTENCY_PREFIX = "lease-reservation:"; /** * Configuration options for the SchematicClient @@ -58,6 +84,15 @@ export interface SchematicOptions { offline?: boolean; /** The default maximum time to wait for a response in milliseconds */ timeoutMs?: number; + /** + * Enable client-side credit lease + reservation behavior on `check` / + * `trackWithReservation`. Omit to keep the SDK lease-unaware. + * + * Lease-bearing checks require DataStream (or replicator mode) so the + * SDK has access to the cached flag + company state needed for local + * gating against the lease balance. + */ + creditLeases?: CreditLeaseConfig; } export interface CheckFlagOptions { @@ -65,6 +100,17 @@ export interface CheckFlagOptions { defaultValue?: boolean | (() => boolean); /** The maximum time to wait for a response in milliseconds */ timeoutMs?: number; + /** + * Preflight inputs forwarded to the WASM rules engine's + * `checkFlagWithOptions`. The engine picks the most specific knob for each + * condition it evaluates. + */ + /** Pre-computed per-credit-id cost. Highest precedence for credit-balance gates. */ + creditCost?: Record; + /** Single integer quantity applied to whatever numeric condition is being evaluated. */ + usage?: number; + /** Per-event-subtype simulated quantity. Preferred when the subtype is known. */ + eventUsage?: Record; } /** @@ -89,6 +135,14 @@ export interface TrackOptions { export interface IdentifyOptions { /** Client-supplied dedupe key. Duplicate events with the same key (scoped to the environment) are dropped server-side for 24 hours. */ idempotencyKey?: string; + /** + * Credit type IDs to acquire leases for in the background after the + * identify event is enqueued. Fire-and-forget — failures are logged but + * never surface to the caller. Equivalent to calling `client.prewarm()` + * directly with the same evalCtx and creditTypeIds. No-op unless + * `creditLeases` is configured on the client. + */ + prewarm?: string[]; } export interface CheckFlagWithEntitlementResponse { @@ -111,6 +165,10 @@ export class SchematicClient extends BaseClient { private flagDefaults: { [key: string]: boolean }; private logger: Logger; private offline: boolean; + private creditLeaseManager?: CreditLeaseManager; + private leaseStore?: ILeaseStore; + private reservations?: IReservationStore; + private prewarmResolveTimeoutMs: number = DEFAULT_PREWARM_RESOLVE_TIMEOUT_MS; /** * Creates a new instance of the SchematicClient @@ -230,6 +288,57 @@ export class SchematicClient extends BaseClient { } } + + // Set up credit lease + reservation plumbing if the caller opted in. + if (opts?.creditLeases && !offline) { + const sweepMs = opts.creditLeases.sweepIntervalMs ?? DEFAULT_SWEEP_INTERVAL_MS; + // Lease + reservation state belongs in a shared cache so gating + // holds across horizontally-scaled pods. Prefer an explicit + // `creditLeases.redisClient`, but otherwise reuse the Redis client + // the DataStream cache is already configured with — so an existing + // Redis setup backs leases automatically, with no second client to + // wire up. Same for the key prefix. + const redisClient = opts.creditLeases.redisClient ?? opts.dataStream?.redisClient; + const keyPrefix = opts.creditLeases.redisKeyPrefix ?? opts.dataStream?.redisKeyPrefix; + if (redisClient) { + // Shared-state backend: lease balance + reservation table live + // in Redis. Lua-script-driven atomicity gives cross-pod gating + // without a separate lock service. + this.leaseStore = new RedisLeaseStore({ + client: redisClient, + keyPrefix, + defaultLeaseDurationMs: opts.creditLeases.defaultLeaseDuration, + }); + this.reservations = new RedisReservationStore({ + client: redisClient, + leaseStore: this.leaseStore, + sweepIntervalMs: sweepMs, + keyPrefix, + }); + } else { + // No shared backend configured: fall back to per-process + // in-memory stores. In a horizontally-scaled deployment each pod + // then acquires and gates against its own leases, which defeats + // the cross-pod over-spend protection that is the point of + // leasing — so warn rather than degrade silently. + logger.warn( + "creditLeases is enabled without a shared Redis backend; lease and reservation " + + "state will be kept per-process. Configure dataStream.redisClient (or " + + "creditLeases.redisClient) so leases gate correctly across multiple SDK instances.", + ); + this.leaseStore = new LeaseStore(); + this.reservations = new ReservationStore(this.leaseStore, sweepMs); + } + this.reservations.startSweep(); + this.creditLeaseManager = new CreditLeaseManager({ + creditsClient: this.credits, + leaseStore: this.leaseStore, + logger, + config: opts.creditLeases, + }); + this.prewarmResolveTimeoutMs = + opts.creditLeases.prewarmResolveTimeoutMs ?? DEFAULT_PREWARM_RESOLVE_TIMEOUT_MS; + } } /** @@ -566,10 +675,23 @@ export class SchematicClient extends BaseClient { } /** - * Gracefully shuts down the client by stopping the event buffer and DataStream client + * Gracefully shuts down the client by stopping the event buffer, the + * reservation sweeper, and the DataStream client. + * + * Outstanding credit leases are deliberately *not* released here. A lease + * is shared across every SDK instance pointed at the same backend (one row + * per company+credit), so a single pod shutting down must not release a + * lease that sibling pods are still drawing on — doing so would refund the + * grant server-side and pull the balance out from under them. Instead we + * let leases reclaim themselves: they expire (client- and server-side) or + * are fully consumed. A caller can't reliably know when a shared lease is + * safe to release, so close() leaves that to expiry. * @returns Promise that resolves when everything has been stopped */ async close(): Promise { + if (this.reservations) { + this.reservations.stop(); + } if (this.datastreamClient) { this.datastreamClient.close(); } @@ -577,11 +699,303 @@ export class SchematicClient extends BaseClient { } /** - * Send a non-blocking event to create or update companies and users - * @param body - The identify event payload containing user properties - * @param options - Optional event metadata (e.g. idempotencyKey) - * @returns Promise that resolves when the event has been enqueued - * @throws Will log error if event enqueueing fails + * Pre-warm a credit lease for each given credit type ID, so the first + * `check()` against it doesn't pay the acquire round-trip. Fire-and-forget + * — failures are logged but don't throw. + * + * When `evalCtx.company` carries only secondary keys (no `id`), `prewarm` + * actively fetches the company over the datastream (waiting up to + * `creditLeases.prewarmResolveTimeoutMs` for the WS to connect), which both + * resolves the id and warms the cache so the first `check()` hits the lease + * path. Covers a brand-new company too: the fetch retries until the server + * has ingested the preceding `identify` and can stream it back. + */ + async prewarm(evalCtx: api.CheckFlagRequestBody, creditTypeIds: string[]): Promise { + if (!this.creditLeaseManager || !this.leaseStore) { + this.logger.debug("prewarm called but creditLeases is not configured"); + return; + } + if (!evalCtx.company || Object.keys(evalCtx.company).length === 0) { + this.logger.debug("prewarm requires a company on evalCtx"); + return; + } + const companyId = await this.resolveCompanyIdWithWait(evalCtx); + if (!companyId) { + this.logger.debug( + `prewarm: company not resolved within ${this.prewarmResolveTimeoutMs}ms for keys ${JSON.stringify(evalCtx.company)} (first check() will acquire)`, + ); + return; + } + await Promise.all( + creditTypeIds.map((id) => + this.creditLeaseManager! + .acquireIfNeeded(companyId, id) + .catch((err) => + this.logger.warn(`prewarm: failed to acquire lease for ${id}: ${err}`), + ), + ), + ); + } + + /** + * Like `resolveCompanyId` but actively fetches the company over the + * datastream when only secondary keys are supplied, warming the cache as a + * side effect. Returns the resolved id, or undefined if the company never + * surfaced within `prewarmResolveTimeoutMs`. + * + * `identify` does not push a company into the datastream cache — companies + * are only streamed in response to a request. So we call `getCompany` + * (cache-first, then sends the request and awaits the push) rather than + * passively polling `getCachedCompany`, which would watch an empty cache + * until it times out. Fetching also primes the cache so the first real + * `check()` hits the lease path instead of falling back. + */ + private async resolveCompanyIdWithWait( + evalCtx: api.CheckFlagRequestBody, + ): Promise { + if (!evalCtx.company) return undefined; + if (evalCtx.company.id) return evalCtx.company.id; + const datastream = this.datastreamClient; + if (!datastream || this.prewarmResolveTimeoutMs <= 0) { + return this.resolveCompanyId(evalCtx); + } + const company = evalCtx.company; + // Fast path — already cached by an earlier check or prewarm. + const cached = await datastream.getCachedCompany(company); + if (cached?.id) return cached.id; + // Actively fetch. Retry across the brief WS-connecting window at boot + // (getCompany throws "not connected" until the socket is ready), + // bounded by prewarmResolveTimeoutMs. + const deadline = Date.now() + this.prewarmResolveTimeoutMs; + for (;;) { + try { + const resolved = await datastream.getCompany(company); + if (resolved?.id) return resolved.id; + } catch (err) { + this.logger.debug(`prewarm: datastream company fetch failed (${err})`); + } + if (Date.now() >= deadline) return undefined; + await new Promise((r) => setTimeout(r, DEFAULT_PREWARM_POLL_INTERVAL_MS)); + } + } + + /** + * Lease-aware feature check. When `creditLeases` is configured and the + * caller passes `usage` / `eventUsage`, this acquires a lease (if needed), + * gates the check against the lease's local balance via WASM, and returns + * a reservation handle on success — pass that handle to + * `trackWithReservation` when the work completes. + * + * When `creditLeases` is not configured (or `usage`/`eventUsage` is omitted) + * this falls through to a plain flag check and returns `{allowed: value}` + * with no reservation, byte-compatible with `checkFlag` semantics. + */ + async check( + key: string, + evalCtx: api.CheckFlagRequestBody, + options?: CheckOptions, + ): Promise { + const fallback = async (): Promise => { + const resp = await this.checkFlagWithEntitlement(evalCtx, key, { + defaultValue: options?.defaultValue, + timeoutMs: options?.timeoutMs, + }); + return { + allowed: resp.value, + value: resp.value, + reason: resp.reason, + entitlement: resp.entitlement, + flagKey: resp.flagKey, + flagId: resp.flagId, + err: resp.err, + }; + }; + + const hasPreflight = + options?.usage !== undefined || + (options?.eventUsage && Object.keys(options.eventUsage).length > 0); + if ( + !hasPreflight || + !this.creditLeaseManager || + !this.leaseStore || + !this.reservations + ) { + return fallback(); + } + + return checkWithLease( + { + leaseStore: this.leaseStore, + reservations: this.reservations, + manager: this.creditLeaseManager, + datastream: this.datastreamClient, + logger: this.logger, + }, + key, + evalCtx, + options, + fallback, + ); + } + + /** + * Consume a reservation issued by `check()`. Refunds the unused slice + * (`creditsReserved - actualQuantity * consumptionRate`) back to the + * lease's local balance and enqueues a Track event with + * `quantity = actualQuantity`. The server-side event processor consumes + * `actualQuantity × consumption_rate` from the company's real credit + * balance. + * + * If the work outlived the reservation's TTL and the sweeper already + * returned the hold to the lease, the local refund has happened but the + * usage must still be billed — so we emit the Track anyway (a "recovery" + * emit). The server bills it against the lease's sub-ledger if the lease is + * still live, or falls through to a direct grant decrement if the server + * lease has itself expired. + * + * Double-billing is prevented server-side: the Track carries a deterministic + * `idempotencyKey` derived from the reservation id, and the events pipeline + * drops duplicates by `(account, env, event_type, key)` for 24h before any + * credit consumption runs. So a recovery emit racing the normal emit, or an + * accidental second `trackWithReservation`, collapses to a single billed + * event — across pods and process restarts, not just within one process. + */ + async trackWithReservation( + reservation: Reservation, + actualQuantity: number, + options?: TrackWithReservationOptions, + ): Promise { + if (this.offline) return; + if (!this.reservations) { + this.logger.warn( + "trackWithReservation called but creditLeases is not configured — emitting unreserved track", + ); + await this.track({ + event: reservation.eventSubtype, + quantity: actualQuantity, + company: reservation.evalCtx.company, + user: reservation.evalCtx.user, + traits: options?.traits, + }); + return; + } + const { track, settledLocally } = await consumeReservationAndBuildEvent( + this.reservations, + reservation, + actualQuantity, + options, + ); + if (!settledLocally) { + this.logger.debug( + `trackWithReservation: reservation ${reservation.id} was no longer in the store (expired/swept or already settled) — emitting recovery track keyed for idempotent server-side dedupe`, + ); + } + // Deterministic idempotency key so duplicate/recovery emits dedupe + // server-side rather than double-billing the lease's sub-ledger. + await this.track(track, { + idempotencyKey: `${RESERVATION_TRACK_IDEMPOTENCY_PREFIX}${reservation.id}`, + }); + } + + /** + * Lease-aware credit balance for display, in the same three-way split the + * server exposes on `listCompanyCreditBalances`, so client and server speak + * one vocabulary: + * + * - `remaining`: credits available to draw right now — the pool a *new* + * lease hold can take. Equals the server balance + * (`company.creditBalances[creditId]`, already net of any active lease + * tranche): `B − L`. + * - `reserved`: the active lease's unsettled hold — credits pulled out of + * the company balance into the lease but not yet settled by a Track. + * This is the lease's unspent local balance plus its open reservations + * (`L − spent`); the client sub-divides the hold into unspent remainder + * vs in-flight reservations, but the server only models the aggregate, + * so we report the aggregate. + * - `settled`: the economic balance the end user should see — + * `remaining + reserved = B − spent`. Invariant to in-flight holds: + * opening a reservation moves credits from `remaining`/lease-remainder + * into the reservation but leaves `settled` unchanged, so a counter + * bound to `settled` tracks settled spend and doesn't dip mid-call. + * + * Bind a user-facing "credits remaining" counter to `settled`. The server + * balance alone reads low — or 0 — while a lease still has unspent local + * balance (a browser SDK that sees only `remaining` shows "0 remaining" + * next to a still-working action); `settled` recovers the true figure. + * + * `reserved` is 0 when no lease is held — or when the held lease has + * expired: the server sweeps expired leases and refunds the unspent hold + * back into the balance, so `remaining` already reflects it. A stale local + * entry past its `expiresAt` is therefore ignored; counting it too would + * double-count. Returns zeros when credit leases / datastream aren't + * configured or the company can't be resolved. + * + * The `reserved` term comes from this process's lease + reservation stores, + * so it (and therefore `settled`) is per-process unless Redis-backed stores + * are configured. + */ + async getCreditBalance( + company: Record, + creditId: string, + ): Promise<{ remaining: number; reserved: number; settled: number }> { + const empty = { remaining: 0, reserved: 0, settled: 0 }; + const datastream = this.datastreamClient; + if (!datastream || !company || Object.keys(company).length === 0) { + return empty; + } + + let snapshot = await datastream.getCachedCompany(company); + if (!snapshot) { + // Not cached yet (e.g. first paint before any check/prewarm). + // Actively fetch so the counter populates; tolerate a cold/closed + // datastream by returning zeros rather than throwing. + try { + snapshot = await datastream.getCompany(company); + } catch (err) { + this.logger.debug(`getCreditBalance: company fetch failed (${err})`); + return empty; + } + } + if (!snapshot?.id) return empty; + + // `remaining` = the drawable pool = the server balance, already net of + // any active lease tranche (`B − L`). + const remaining = snapshot.creditBalances?.[creditId] ?? 0; + const leaseEntry = await this.leaseStore?.get(snapshot.id, creditId); + // Only count a live lease. Past expiresAt the server has (or imminently + // will) sweep and refund the hold, so `remaining` already accounts for it. + const leaseLive = + leaseEntry !== undefined && + new Date(leaseEntry.expiresAt).getTime() > Date.now(); + // `reserved` = the lease's whole unsettled hold (`L − spent`): its + // unspent local balance plus its open reservations. Both are 0 without + // a live lease. + const leaseRemainder = leaseLive ? leaseEntry.localRemainingCredits : 0; + const openReservations = leaseLive + ? ((await this.reservations?.reservedCredits(snapshot.id, creditId)) ?? 0) + : 0; + const reserved = leaseRemainder + openReservations; + + return { remaining, reserved, settled: remaining + reserved }; + } + + private async resolveCompanyId(evalCtx: api.CheckFlagRequestBody): Promise { + if (!evalCtx.company) return undefined; + // If the caller passed `id`, use that directly. + if (evalCtx.company.id) return evalCtx.company.id; + // Otherwise, the datastream cache can resolve secondary keys → id. + if (this.datastreamClient) { + const cached = await this.datastreamClient.getCachedCompany(evalCtx.company); + return cached?.id; + } + return undefined; + } + + /** + * Send a non-blocking event to create or update companies and users. + * Pass `options.prewarm` with credit type IDs to kick off lease acquires in + * the background after the identify event is enqueued (no-op unless + * `creditLeases` is configured on the client). */ async identify(body: api.EventBodyIdentify, options?: IdentifyOptions): Promise { if (this.offline) return; @@ -591,6 +1005,23 @@ export class SchematicClient extends BaseClient { } catch (err) { this.logger.error(`Error sending identify event: ${err}`); } + + if (options?.prewarm && options.prewarm.length > 0) { + const evalCtx: api.CheckFlagRequestBody = { + company: body.company?.keys, + user: body.keys, + }; + // Force-flush so the server processes the identify ASAP — without + // it, the company may sit in the local buffer for up to the + // buffer's flush interval before the server even sees it, and + // prewarm's bounded poll would just be waiting on us. Fire-and-forget. + void this.eventBuffer.flush().catch((err) => { + this.logger.debug(`identify flush before prewarm failed: ${err}`); + }); + void this.prewarm(evalCtx, options.prewarm).catch((err) => { + this.logger.warn(`identify prewarm failed: ${err}`); + }); + } } /** diff --git a/tests/unit/credits/check-and-track.test.ts b/tests/unit/credits/check-and-track.test.ts new file mode 100644 index 00000000..459923a6 --- /dev/null +++ b/tests/unit/credits/check-and-track.test.ts @@ -0,0 +1,544 @@ +import { SchematicClient } from "../../../src/wrapper"; + +const mockCheckFlag = jest.fn(); +const mockAcquireCreditLease = jest.fn(); +const mockExtendCreditLease = jest.fn(); +const mockReleaseCreditLease = jest.fn(); + +jest.mock("../../../src/Client", () => { + class MockBaseClient { + features = { + checkFlag: mockCheckFlag, + checkFlags: jest.fn().mockResolvedValue({ data: { flags: [] } }), + }; + credits = { + acquireCreditLease: mockAcquireCreditLease, + extendCreditLease: mockExtendCreditLease, + releaseCreditLease: mockReleaseCreditLease, + }; + events = {}; + } + return { SchematicClient: MockBaseClient }; +}); + +const mockEventBufferPush = jest.fn(); +jest.mock("../../../src/events", () => ({ + EventBuffer: jest.fn().mockImplementation(() => ({ + push: mockEventBufferPush, + flush: jest.fn().mockResolvedValue(undefined), + stop: jest.fn().mockResolvedValue(undefined), + })), +})); + +// Stub DataStreamClient to a fixed shape — we manually wire its methods so +// the lease path has everything it needs without spinning up a websocket. +const mockDataStream = { + start: jest.fn().mockResolvedValue(undefined), + close: jest.fn(), + isConnected: jest.fn().mockReturnValue(true), + checkFlag: jest.fn(), + updateCompanyMetrics: jest.fn().mockResolvedValue(undefined), + getFlag: jest.fn(), + getCachedCompany: jest.fn(), + getCompany: jest.fn(), + getCachedUser: jest.fn().mockResolvedValue(null), + getRulesEngine: jest.fn(), +}; + +jest.mock("../../../src/datastream", () => ({ + DataStreamClient: jest.fn().mockImplementation(() => mockDataStream), +})); + +// Stub rules engine +const mockRulesEngine = { + initialize: jest.fn().mockResolvedValue(undefined), + isInitialized: jest.fn().mockReturnValue(true), + checkFlag: jest.fn(), + checkFlagWithOptions: jest.fn(), + getVersionKey: jest.fn().mockReturnValue("1"), +}; +jest.mock("../../../src/rules-engine", () => ({ + RulesEngineClient: jest.fn().mockImplementation(() => mockRulesEngine), +})); + +mockDataStream.getRulesEngine.mockReturnValue(mockRulesEngine); + +const flag = { + accountId: "acct", + defaultValue: false, + environmentId: "env", + id: "flag_1", + key: "inference", + rules: [ + { + accountId: "acct", + conditionGroups: [], + conditions: [ + { + accountId: "acct", + conditionType: "credit", + consumptionRate: 10, + creditId: "bilcr_inference", + environmentId: "env", + eventSubtype: "inference_tokens", + id: "cond_1", + operator: "gte", + resourceIds: [], + traitValue: "0", + }, + ], + environmentId: "env", + id: "rule_1", + name: "credit gate", + priority: 1, + ruleType: "standard", + value: true, + }, + ], +}; + +// Build a single-condition credit rule for a given credit type. Used to +// assemble flags that meter one feature across multiple credit types. +function makeCreditRule(id: string, creditId: string, consumptionRate: number) { + return { + accountId: "acct", + conditionGroups: [], + conditions: [ + { + accountId: "acct", + conditionType: "credit", + consumptionRate, + creditId, + environmentId: "env", + eventSubtype: "inference_tokens", + id: `cond_${creditId}`, + operator: "gte", + resourceIds: [], + traitValue: "0", + }, + ], + environmentId: "env", + id, + name: id, + priority: 1, + ruleType: "standard", + value: true, + }; +} + +const company = { + id: "co_1", + accountId: "acct", + environmentId: "env", + keys: { id: "co_1" }, + metrics: [], + creditBalances: { bilcr_inference: 5000 }, + planIds: [], + planVersionIds: [], + billingProductIds: [], + subscription: null, + traits: [], + rules: [], +}; + +function configureSuccessfulAcquire() { + mockAcquireCreditLease.mockResolvedValue({ + data: { + id: "lse_1", + companyId: "co_1", + creditTypeId: "bilcr_inference", + grantedAmount: 1000, + expiresAt: new Date(Date.now() + 5 * 60_000), + createdAt: new Date(), + updatedAt: new Date(), + }, + params: {}, + }); +} + +function configureFailingAcquire() { + mockAcquireCreditLease.mockRejectedValue(new Error("lease 503")); +} + +function configureDataStream() { + mockDataStream.getFlag.mockResolvedValue(flag); + mockDataStream.getCachedCompany.mockResolvedValue(company); +} + +function makeClient() { + return new SchematicClient({ + apiKey: "test-key", + useDataStream: true, + creditLeases: { + defaultLeaseDuration: 5 * 60_000, + defaultReservationTTL: 60_000, + defaultLeaseSize: 1000, + lowWaterMark: 0.25, + sweepIntervalMs: 60_000, + }, + logger: { + debug: jest.fn(), + info: jest.fn(), + warn: jest.fn(), + error: jest.fn(), + }, + }); +} + +beforeEach(() => { + jest.clearAllMocks(); + mockDataStream.getRulesEngine.mockReturnValue(mockRulesEngine); + mockDataStream.isConnected.mockReturnValue(true); +}); + +describe("client.check (lease path)", () => { + it("issues a reservation when the engine says allowed", async () => { + configureSuccessfulAcquire(); + configureDataStream(); + mockRulesEngine.checkFlagWithOptions.mockResolvedValue({ + value: true, + reason: "matched", + flagKey: "inference", + flagId: "flag_1", + entitlement: { + featureId: "feat", + featureKey: "inference", + valueType: "credit_burndown", + creditId: "bilcr_inference", + creditTotal: 1000, + creditUsed: 0, + creditRemaining: 1000, + }, + }); + + const client = makeClient(); + const result = await client.check( + "inference", + { company: { id: "co_1" } }, + { eventUsage: { inference_tokens: 50 } }, + ); + + expect(result.allowed).toBe(true); + expect(result.reservation).toBeDefined(); + expect(result.reservation?.creditsReserved).toBe(500); + expect(result.reservation?.consumptionRate).toBe(10); + expect(result.reservation?.quantityReserved).toBe(50); + expect(mockAcquireCreditLease).toHaveBeenCalledTimes(1); + // Substituted balance flows into WASM: pre-reservation localRemaining = 1000. + const callArgs = mockRulesEngine.checkFlagWithOptions.mock.calls[0]; + expect(callArgs[1].creditBalances.bilcr_inference).toBe(1000); + expect(callArgs[3]).toEqual({ eventUsage: { inference_tokens: 50 } }); + await client.close(); + }); + + it("leases the credit the company's matched plan uses when a flag mixes credit types", async () => { + // `inference` is entitled via two plans: a legacy USD-cents credit + // (declared first on the flag) and an AI-credits credit. The company is + // on the AI-credits plan, so the lease must target AI credits even + // though the USD condition appears first. The engine probe reports the + // matched plan's credit; we lease that, not the first-declared one. + const mixedCreditFlag = { + ...flag, + rules: [ + makeCreditRule("rule_usd", "bilcr_usd", 10), + makeCreditRule("rule_ai", "bilcr_ai", 5), + ], + }; + mockDataStream.getFlag.mockResolvedValue(mixedCreditFlag); + mockDataStream.getCachedCompany.mockResolvedValue({ + ...company, + creditBalances: { bilcr_usd: 0, bilcr_ai: 5000 }, + }); + mockAcquireCreditLease.mockResolvedValue({ + data: { + id: "lse_ai", + companyId: "co_1", + creditTypeId: "bilcr_ai", + grantedAmount: 1000, + expiresAt: new Date(Date.now() + 5 * 60_000), + createdAt: new Date(), + updatedAt: new Date(), + }, + params: {}, + }); + // Probe + final eval both report the matched plan's credit (AI credits). + mockRulesEngine.checkFlagWithOptions.mockResolvedValue({ + value: true, + reason: "matched ai plan", + flagKey: "inference", + flagId: "flag_1", + entitlement: { + featureId: "feat", + featureKey: "inference", + valueType: "credit_burndown", + creditId: "bilcr_ai", + creditRemaining: 5000, + }, + }); + + const client = makeClient(); + const result = await client.check( + "inference", + { company: { id: "co_1" } }, + { eventUsage: { inference_tokens: 50 } }, + ); + + expect(result.allowed).toBe(true); + // Leased AI credits (the matched plan), not the first-declared USD credit. + expect(result.reservation?.creditTypeId).toBe("bilcr_ai"); + expect(result.reservation?.consumptionRate).toBe(5); + expect(result.reservation?.creditsReserved).toBe(250); + expect(mockAcquireCreditLease).toHaveBeenCalledTimes(1); + expect(mockAcquireCreditLease.mock.calls[0][0].creditTypeId).toBe("bilcr_ai"); + // One probe eval (raw balance) + one gating eval (substituted balance). + expect(mockRulesEngine.checkFlagWithOptions).toHaveBeenCalledTimes(2); + await client.close(); + }); + + it("returns allowed=false and refunds reservation when engine denies", async () => { + configureSuccessfulAcquire(); + configureDataStream(); + mockRulesEngine.checkFlagWithOptions.mockResolvedValue({ + value: false, + reason: "denied", + flagKey: "inference", + }); + + const client = makeClient(); + const result = await client.check( + "inference", + { company: { id: "co_1" } }, + { eventUsage: { inference_tokens: 50 } }, + ); + + expect(result.allowed).toBe(false); + expect(result.reservation).toBeUndefined(); + await client.close(); + }); + + it("fails closed when lease acquire errors and onAcquireFailure='fail-closed'", async () => { + configureFailingAcquire(); + configureDataStream(); + + const client = makeClient(); + const result = await client.check( + "inference", + { company: { id: "co_1" } }, + { eventUsage: { inference_tokens: 50 }, onAcquireFailure: "fail-closed" }, + ); + + expect(result.allowed).toBe(false); + expect(result.reservation).toBeUndefined(); + expect(mockRulesEngine.checkFlagWithOptions).not.toHaveBeenCalled(); + await client.close(); + }); + + it("defaults to fail-closed when onAcquireFailure is not specified", async () => { + configureFailingAcquire(); + configureDataStream(); + + const client = makeClient(); + const result = await client.check( + "inference", + { company: { id: "co_1" } }, + { eventUsage: { inference_tokens: 50 } }, + ); + + expect(result.allowed).toBe(false); + expect(result.reservation).toBeUndefined(); + expect(mockRulesEngine.checkFlagWithOptions).not.toHaveBeenCalled(); + await client.close(); + }); + + it("respects explicit fail-open override on acquire failure", async () => { + configureFailingAcquire(); + configureDataStream(); + + const client = makeClient(); + const result = await client.check( + "inference", + { company: { id: "co_1" } }, + { eventUsage: { inference_tokens: 50 }, onAcquireFailure: "fail-open" }, + ); + + expect(result.allowed).toBe(true); + expect(result.reservation).toBeUndefined(); + expect(mockRulesEngine.checkFlagWithOptions).not.toHaveBeenCalled(); + await client.close(); + }); + + it("falls back to plain check when no preflight options are provided", async () => { + configureDataStream(); + mockCheckFlag.mockResolvedValue({ + data: { value: true, flag: "inference", reason: "match" }, + }); + + const client = makeClient(); + const result = await client.check("inference", { company: { id: "co_1" } }); + + expect(result.allowed).toBe(true); + expect(mockAcquireCreditLease).not.toHaveBeenCalled(); + await client.close(); + }); +}); + +describe("client.prewarm (datastream resolution)", () => { + it("actively fetches a not-yet-cached company over the datastream before acquiring", async () => { + configureSuccessfulAcquire(); + mockDataStream.getFlag.mockResolvedValue(flag); + // Not cached yet. prewarm actively fetches via getCompany (identify + // doesn't push the company into the cache on its own). Reject once to + // exercise the WS-connecting retry, then surface the company. + mockDataStream.getCachedCompany.mockResolvedValue(null); + mockDataStream.getCompany + .mockRejectedValueOnce(new Error("DataStream client is not connected")) + .mockResolvedValue(company); + + const client = makeClient(); + await client.prewarm( + { company: { external_id: "ext-co-1" } }, + ["bilcr_inference"], + ); + + expect(mockDataStream.getCompany).toHaveBeenCalledWith({ external_id: "ext-co-1" }); + expect(mockAcquireCreditLease).toHaveBeenCalledTimes(1); + expect(mockAcquireCreditLease).toHaveBeenCalledWith( + expect.objectContaining({ companyId: "co_1", creditTypeId: "bilcr_inference" }), + ); + await client.close(); + }); + + it("gives up after prewarmResolveTimeoutMs without acquiring", async () => { + configureSuccessfulAcquire(); + mockDataStream.getFlag.mockResolvedValue(flag); + mockDataStream.getCachedCompany.mockResolvedValue(null); + // Company never surfaces over the datastream. + mockDataStream.getCompany.mockRejectedValue( + new Error("DataStream client is not connected"), + ); + + const client = new SchematicClient({ + apiKey: "test-key", + useDataStream: true, + creditLeases: { + defaultLeaseSize: 1000, + sweepIntervalMs: 60_000, + prewarmResolveTimeoutMs: 250, + }, + logger: { debug: jest.fn(), info: jest.fn(), warn: jest.fn(), error: jest.fn() }, + }); + await client.prewarm( + { company: { external_id: "ext-co-missing" } }, + ["bilcr_inference"], + ); + + expect(mockAcquireCreditLease).not.toHaveBeenCalled(); + await client.close(); + }); +}); + +describe("client.trackWithReservation", () => { + it("refunds unused credits and emits a Track event with actual quantity", async () => { + configureSuccessfulAcquire(); + configureDataStream(); + mockRulesEngine.checkFlagWithOptions.mockResolvedValue({ + value: true, + reason: "matched", + flagKey: "inference", + flagId: "flag_1", + }); + + const client = makeClient(); + const res = await client.check( + "inference", + { company: { id: "co_1" } }, + { eventUsage: { inference_tokens: 50 } }, + ); + if (!res.reservation) throw new Error("expected reservation"); + + // Local balance after reservation: 1000 - 500 = 500 + await client.trackWithReservation(res.reservation, 20); + + // Track event was enqueued with actualQuantity + const pushed = mockEventBufferPush.mock.calls.find( + (call) => call[0]?.eventType === "track", + ); + expect(pushed).toBeDefined(); + expect(pushed?.[0].body.event).toBe("inference_tokens"); + expect(pushed?.[0].body.quantity).toBe(20); + expect(pushed?.[0].body.company).toEqual({ id: "co_1" }); + await client.close(); + }); + + it("emits a recovery Track keyed for idempotent dedupe when the reservation is gone before track", async () => { + configureSuccessfulAcquire(); + configureDataStream(); + const client = new SchematicClient({ + apiKey: "test-key", + useDataStream: true, + creditLeases: { defaultLeaseSize: 1000, sweepIntervalMs: 60_000 }, + logger: { debug: jest.fn(), info: jest.fn(), warn: jest.fn(), error: jest.fn() }, + }); + const reservation = { + id: "missing", + leaseId: "lse_x", + companyId: "co_1", + creditTypeId: "bilcr_inference", + eventSubtype: "inference_tokens", + quantityReserved: 10, + creditsReserved: 100, + consumptionRate: 10, + expiresAt: new Date(Date.now() + 60_000), + evalCtx: { company: { id: "co_1" } }, + }; + + // Reservation is gone from the store (swept after TTL), but the work + // completed — we still owe the server a Track so the usage is billed. + await client.trackWithReservation(reservation, 5); + const trackCalls = () => + mockEventBufferPush.mock.calls.filter((call) => call[0]?.eventType === "track"); + expect(trackCalls()).toHaveLength(1); + const recovery = trackCalls()[0][0]; + expect(recovery.body.event).toBe("inference_tokens"); + expect(recovery.body.quantity).toBe(5); + expect(recovery.body.leaseId).toBe("lse_x"); + // The Track carries a deterministic idempotency key derived from the + // reservation id, so the server dedupes duplicate/recovery emits. + expect(recovery.idempotencyKey).toBe("lease-reservation:missing"); + + // A second call re-emits with the SAME key — the SDK relies on the + // server's (account, env, event_type, key) dedupe to collapse them to + // one billed event, rather than suppressing client-side (which wouldn't + // hold across pods or restarts). + await client.trackWithReservation(reservation, 5); + const calls = trackCalls(); + expect(calls).toHaveLength(2); + expect(calls[1][0].idempotencyKey).toBe("lease-reservation:missing"); + await client.close(); + }); + + it("keys the normal settle Track with the reservation id too", async () => { + configureSuccessfulAcquire(); + configureDataStream(); + mockRulesEngine.checkFlagWithOptions.mockResolvedValue({ + value: true, + reason: "matched", + flagKey: "inference", + flagId: "flag_1", + }); + const client = makeClient(); + const res = await client.check( + "inference", + { company: { id: "co_1" } }, + { eventUsage: { inference_tokens: 50 } }, + ); + if (!res.reservation) throw new Error("expected reservation"); + + await client.trackWithReservation(res.reservation, 20); + const pushed = mockEventBufferPush.mock.calls.find( + (call) => call[0]?.eventType === "track", + ); + expect(pushed?.[0].idempotencyKey).toBe(`lease-reservation:${res.reservation.id}`); + await client.close(); + }); +}); diff --git a/tests/unit/credits/fake-redis.ts b/tests/unit/credits/fake-redis.ts new file mode 100644 index 00000000..863a201a --- /dev/null +++ b/tests/unit/credits/fake-redis.ts @@ -0,0 +1,240 @@ +import type { RedisClient } from "../../../src/cache/redis"; + +/** + * In-memory implementation of the subset of the `RedisClient` interface used + * by `RedisLeaseStore` + `RedisReservationStore`. The `eval` method + * interprets the specific Lua scripts those stores send — it's not a generic + * Lua engine, just enough to exercise the atomic semantics in tests. + * + * Atomicity: the JS event loop is single-threaded and our Lua paths are all + * synchronous against the in-memory maps, so calling them from concurrent + * promises emulates real Redis Lua atomicity faithfully enough for these + * tests. + */ +export function makeFakeRedis(): RedisClient { + const strings = new Map(); + const hashes = new Map>(); + const sets = new Map>(); + const zsets = new Map>(); // member -> score + const expirations = new Map(); // key -> epoch ms + + const checkExpiry = (key: string): void => { + const exp = expirations.get(key); + if (exp !== undefined && exp <= Date.now()) { + strings.delete(key); + hashes.delete(key); + sets.delete(key); + zsets.delete(key); + expirations.delete(key); + } + }; + + const hgetAll = (key: string): Record => { + checkExpiry(key); + const h = hashes.get(key); + if (!h) return {}; + return Object.fromEntries(h.entries()); + }; + + const hset = (key: string, field: string, value: string): void => { + if (!hashes.has(key)) hashes.set(key, new Map()); + hashes.get(key)!.set(field, value); + }; + + const hget = (key: string, field: string): string | null => { + checkExpiry(key); + return hashes.get(key)?.get(field) ?? null; + }; + + const del = (key: string): void => { + strings.delete(key); + hashes.delete(key); + sets.delete(key); + zsets.delete(key); + expirations.delete(key); + }; + + const zadd = (key: string, score: number, member: string): void => { + if (!zsets.has(key)) zsets.set(key, new Map()); + zsets.get(key)!.set(member, score); + }; + + const zrem = (key: string, member: string): void => { + zsets.get(key)?.delete(member); + }; + + const interpretScript = (script: string, keys: string[], args: string[]): unknown => { + const trimmed = script.trim(); + // Every script below is single-key (KEYS[1] only) — mirroring the real + // stores, which avoid multi-key Lua so they stay Cluster-safe. + // Match by a stable substring per script. + if (trimmed.includes("local existing_id = redis.call('HGET', KEYS[1], 'leaseId')") && + trimmed.includes("if existing_id and existing_expiry > now then")) { + // REPLACE_SCRIPT — keeps any live lease (regardless of id). + // The real script reads `now` from the Redis server clock + // (redis.call('TIME')); we emulate that with the local clock since + // the fake is in-process. + const [leaseHashKey] = keys; + const [newId, newGranted, newExpiryStr, grace, companyId, creditTypeId] = args; + const newExpiry = Number(newExpiryStr); + const now = Date.now(); + const existingId = hget(leaseHashKey, "leaseId"); + const existingExpiry = Number(hget(leaseHashKey, "expiresAt") ?? "0"); + if (existingId && existingExpiry > now) { + return 0; + } + del(leaseHashKey); + hset(leaseHashKey, "leaseId", newId); + hset(leaseHashKey, "companyId", companyId); + hset(leaseHashKey, "creditTypeId", creditTypeId); + hset(leaseHashKey, "grantedAmount", newGranted); + hset(leaseHashKey, "localRemainingCredits", newGranted); + hset(leaseHashKey, "expiresAt", newExpiryStr); + expirations.set(leaseHashKey, newExpiry + Number(grace)); + return 1; + } + if (trimmed.includes("local raw = redis.call('HGET', KEYS[1], 'localRemainingCredits')") && + trimmed.includes("if remaining < requested then return 0 end")) { + // TRY_RESERVE_SCRIPT — `now` is the Redis server clock in the real + // script (redis.call('TIME')); emulated with the local clock here. + const [leaseHashKey] = keys; + const remainingRaw = hget(leaseHashKey, "localRemainingCredits"); + if (remainingRaw === null) return 0; + const expiry = Number(hget(leaseHashKey, "expiresAt") ?? "0"); + const now = Date.now(); + if (expiry <= now) return 0; + const remaining = Number(remainingRaw); + const requested = Number(args[0]); + if (remaining < requested) return 0; + hset(leaseHashKey, "localRemainingCredits", String(remaining - requested)); + return 1; + } + if (trimmed.includes("local new_balance = remaining + refund") && + trimmed.includes("if new_balance > granted then new_balance = granted end")) { + // REFUND_SCRIPT + const [leaseHashKey] = keys; + const rawRemaining = hget(leaseHashKey, "localRemainingCredits"); + if (rawRemaining === null) return 0; + const remaining = Number(rawRemaining); + const granted = Number(hget(leaseHashKey, "grantedAmount") ?? "0"); + const refund = Number(args[0]); + const newBalance = Math.min(remaining + refund, granted); + hset(leaseHashKey, "localRemainingCredits", String(newBalance)); + return 1; + } + if (trimmed.includes("local raw_granted = redis.call('HGET', KEYS[1], 'grantedAmount')") && + trimmed.includes("'localRemainingCredits', tostring(remaining + add)")) { + // EXTEND_SCRIPT + const [leaseHashKey] = keys; + const rawGranted = hget(leaseHashKey, "grantedAmount"); + if (rawGranted === null) return 0; + const granted = Number(rawGranted); + const remaining = Number(hget(leaseHashKey, "localRemainingCredits") ?? "0"); + const add = Number(args[0]); + const newExpiry = Number(args[1]); + const grace = Number(args[2]); + hset(leaseHashKey, "grantedAmount", String(granted + add)); + hset(leaseHashKey, "localRemainingCredits", String(remaining + add)); + hset(leaseHashKey, "expiresAt", args[1]); + expirations.set(leaseHashKey, newExpiry + grace); + return 1; + } + if (trimmed.includes("local raw = redis.call('HGETALL', KEYS[1])") && + trimmed.includes("redis.call('DEL', KEYS[1])") && + trimmed.includes("return raw")) { + // CLAIM_SCRIPT (reservation) — atomic read + delete, returns the + // flat [field, value, ...] array (or nil when already gone). + const [resHashKey] = keys; + const raw = hgetAll(resHashKey); + if (Object.keys(raw).length === 0) return null; + del(resHashKey); + const flat: string[] = []; + for (const [k, v] of Object.entries(raw)) { + flat.push(k, v); + } + return flat; + } + throw new Error(`fake-redis: unrecognized Lua script:\n${script}`); + }; + + return { + // Basic string ops (unused by lease stores but required by interface) + async get(key) { + checkExpiry(key); + return strings.get(key) ?? null; + }, + async set(key, value) { + strings.set(key, String(value)); + }, + async setEx(key, seconds, value) { + strings.set(key, String(value)); + expirations.set(key, Date.now() + seconds * 1000); + }, + async del(keyOrKeys) { + const arr = Array.isArray(keyOrKeys) ? keyOrKeys : [keyOrKeys]; + for (const k of arr) del(k); + }, + async *scanIterator(_options) { + for (const k of [ + ...strings.keys(), + ...hashes.keys(), + ...sets.keys(), + ...zsets.keys(), + ]) { + yield k; + } + }, + // Hash ops + async hSet(key, field, value) { + if (typeof field === "object" && field !== null) { + for (const [f, v] of Object.entries(field)) { + hset(key, f, String(v)); + } + return; + } + hset(key, field, String(value)); + }, + async hGet(key, field) { + return hget(key, field) ?? undefined; + }, + async hGetAll(key) { + return hgetAll(key); + }, + async hDel(key, field) { + const fields = Array.isArray(field) ? field : [field]; + const h = hashes.get(key); + if (!h) return; + for (const f of fields) h.delete(f); + }, + // Sorted-set ops + async zAdd(key, members) { + const arr = Array.isArray(members) ? members : [members]; + for (const m of arr) zadd(key, m.score, m.value); + }, + async zRangeByScore(key, min, max) { + checkExpiry(key); + const z = zsets.get(key); + if (!z) return []; + const minN = min === "-inf" ? Number.NEGATIVE_INFINITY : Number(min); + const maxN = max === "+inf" ? Number.POSITIVE_INFINITY : Number(max); + return Array.from(z.entries()) + .filter(([, score]) => score >= minN && score <= maxN) + .sort(([, a], [, b]) => a - b) + .map(([m]) => m); + }, + async zRem(key, member) { + const members = Array.isArray(member) ? member : [member]; + for (const m of members) zrem(key, m); + }, + async eval(script, options) { + return interpretScript(script, options.keys, options.arguments); + }, + async pExpireAt(key, timestamp) { + expirations.set(key, timestamp); + }, + async zCard(key) { + checkExpiry(key); + return zsets.get(key)?.size ?? 0; + }, + }; +} diff --git a/tests/unit/credits/lease-manager.test.ts b/tests/unit/credits/lease-manager.test.ts new file mode 100644 index 00000000..a49d9d65 --- /dev/null +++ b/tests/unit/credits/lease-manager.test.ts @@ -0,0 +1,383 @@ +import { CreditLeaseManager } from "../../../src/credits/lease-manager"; +import { LeaseStore } from "../../../src/credits/lease-store"; +import { RedisLeaseStore } from "../../../src/credits/redis-lease-store"; +import type { ILeaseStore } from "../../../src/credits/lease-store"; +import type { Logger } from "../../../src/logger"; +import { makeFakeRedis } from "./fake-redis"; + +function makeLogger(): Logger { + return { + debug: jest.fn(), + info: jest.fn(), + warn: jest.fn(), + error: jest.fn(), + }; +} + +function makeManager(creditsClient: { [k: string]: jest.Mock }) { + const store = new LeaseStore(); + const manager = new CreditLeaseManager({ + // biome-ignore lint/suspicious/noExplicitAny: stubbed client + creditsClient: creditsClient as any, + leaseStore: store, + logger: makeLogger(), + config: { + defaultLeaseDuration: 5 * 60_000, + defaultReservationTTL: 60_000, + defaultLeaseSize: 1000, + lowWaterMark: 0.25, + }, + }); + return { manager, store }; +} + +describe("CreditLeaseManager", () => { + it("acquireIfNeeded calls acquireCreditLease and installs the lease", async () => { + const expiresAt = new Date(Date.now() + 5 * 60_000); + const creditsClient = { + acquireCreditLease: jest.fn().mockResolvedValue({ + data: { + id: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 1000, + expiresAt, + createdAt: new Date(), + updatedAt: new Date(), + }, + params: {}, + }), + extendCreditLease: jest.fn(), + releaseCreditLease: jest.fn(), + }; + const { manager, store } = makeManager(creditsClient); + + const entry = await manager.acquireIfNeeded("co_1", "ct_1"); + expect(creditsClient.acquireCreditLease).toHaveBeenCalledTimes(1); + expect(entry?.leaseId).toBe("lse_1"); + expect(entry?.grantedAmount).toBe(1000); + expect(store.get("co_1", "ct_1")?.localRemainingCredits).toBe(1000); + }); + + it("acquireIfNeeded is single-flight for concurrent callers", async () => { + let resolve!: (v: unknown) => void; + const pending = new Promise((r) => (resolve = r)); + const creditsClient = { + acquireCreditLease: jest.fn().mockReturnValue(pending), + extendCreditLease: jest.fn(), + releaseCreditLease: jest.fn(), + }; + const { manager } = makeManager(creditsClient); + + const p1 = manager.acquireIfNeeded("co_1", "ct_1"); + const p2 = manager.acquireIfNeeded("co_1", "ct_1"); + const p3 = manager.acquireIfNeeded("co_1", "ct_1"); + + resolve({ + data: { + id: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 1000, + expiresAt: new Date(Date.now() + 5 * 60_000), + createdAt: new Date(), + updatedAt: new Date(), + }, + params: {}, + }); + + await Promise.all([p1, p2, p3]); + expect(creditsClient.acquireCreditLease).toHaveBeenCalledTimes(1); + }); + + it("acquireIfNeeded reuses a live lease (no second wire call)", async () => { + const creditsClient = { + acquireCreditLease: jest.fn().mockResolvedValue({ + data: { + id: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 1000, + expiresAt: new Date(Date.now() + 5 * 60_000), + createdAt: new Date(), + updatedAt: new Date(), + }, + params: {}, + }), + extendCreditLease: jest.fn(), + releaseCreditLease: jest.fn(), + }; + const { manager } = makeManager(creditsClient); + await manager.acquireIfNeeded("co_1", "ct_1"); + await manager.acquireIfNeeded("co_1", "ct_1"); + expect(creditsClient.acquireCreditLease).toHaveBeenCalledTimes(1); + }); + + it("acquireIfNeeded re-acquires over an expired slot, replacing it in place", async () => { + // No explicit drop happens inside acquireIfNeeded — `replace` overwrites + // the expired entry atomically. Seed an already-expired lease and verify + // the fresh one supplants it without a redundant-release call. + const creditsClient = { + acquireCreditLease: jest.fn().mockResolvedValue({ + data: { + id: "lse_fresh", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 1000, + expiresAt: new Date(Date.now() + 5 * 60_000), + createdAt: new Date(), + updatedAt: new Date(), + }, + params: {}, + }), + extendCreditLease: jest.fn(), + releaseCreditLease: jest.fn(), + }; + const { manager, store } = makeManager(creditsClient); + // Seed an already-expired lease in the slot. + await store.replace({ + leaseId: "lse_stale", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 1000, + expiresAt: new Date(Date.now() - 1), + }); + + const entry = await manager.acquireIfNeeded("co_1", "ct_1"); + expect(creditsClient.acquireCreditLease).toHaveBeenCalledTimes(1); + // Fresh lease supplants the stale one in place, full balance restored. + expect(entry?.leaseId).toBe("lse_fresh"); + expect(entry?.localRemainingCredits).toBe(1000); + expect(store.get("co_1", "ct_1")?.leaseId).toBe("lse_fresh"); + // The redundant-lease release path must NOT fire: the slot was expired, + // so `replace` wrote (returned true) rather than keeping a live lease. + expect(creditsClient.releaseCreditLease).not.toHaveBeenCalled(); + }); + + it("maybeExtendInBackground triggers extend when below low water mark", async () => { + const creditsClient = { + acquireCreditLease: jest.fn().mockResolvedValue({ + data: { + id: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 1000, + expiresAt: new Date(Date.now() + 5 * 60_000), + createdAt: new Date(), + updatedAt: new Date(), + }, + params: {}, + }), + extendCreditLease: jest.fn().mockResolvedValue({ + data: { + id: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 2000, + expiresAt: new Date(Date.now() + 5 * 60_000), + createdAt: new Date(), + updatedAt: new Date(), + }, + params: {}, + }), + releaseCreditLease: jest.fn(), + }; + const { manager, store } = makeManager(creditsClient); + await manager.acquireIfNeeded("co_1", "ct_1"); + // Spend down to below 25% + await store.tryReserve("co_1", "ct_1", 800); + await manager.maybeExtendInBackground("co_1", "ct_1"); + expect(creditsClient.extendCreditLease).toHaveBeenCalledTimes(1); + expect(store.get("co_1", "ct_1")?.grantedAmount).toBe(2000); + }); + + it("maybeExtendInBackground extends when requiredCredits exceeds local remaining even above watermark", async () => { + const creditsClient = { + acquireCreditLease: jest.fn().mockResolvedValue({ + data: { + id: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 1000, + expiresAt: new Date(Date.now() + 5 * 60_000), + createdAt: new Date(), + updatedAt: new Date(), + }, + params: {}, + }), + extendCreditLease: jest.fn().mockResolvedValue({ + data: { + id: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 2000, + expiresAt: new Date(Date.now() + 5 * 60_000), + createdAt: new Date(), + updatedAt: new Date(), + }, + params: {}, + }), + releaseCreditLease: jest.fn(), + }; + const { manager, store } = makeManager(creditsClient); + await manager.acquireIfNeeded("co_1", "ct_1"); + // Spend a little — still well above the 25% watermark (900/1000 = 90%). + await store.tryReserve("co_1", "ct_1", 100); + // Without the hint, this would no-op (ratio > watermark). + await manager.maybeExtendInBackground("co_1", "ct_1"); + expect(creditsClient.extendCreditLease).not.toHaveBeenCalled(); + // Caller asks for 1500 credits worth — we only have 900 local, so extend. + await manager.maybeExtendInBackground("co_1", "ct_1", 1500); + expect(creditsClient.extendCreditLease).toHaveBeenCalledTimes(1); + expect(store.get("co_1", "ct_1")?.grantedAmount).toBe(2000); + }); + + it("does not conflate concurrent acquire and extend on the same key", async () => { + // Pre-install a live lease with a sub-watermark balance so an extend + // is warranted. We then hold the extend mid-flight and fire an + // acquireIfNeeded against an *expired* slot for a different lease id — + // the two operations must not share inflight state. + let releaseExtend!: (v: unknown) => void; + const extendPending = new Promise((r) => (releaseExtend = r)); + const creditsClient = { + acquireCreditLease: jest.fn().mockResolvedValue({ + data: { + id: "lse_fresh", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 1000, + expiresAt: new Date(Date.now() + 5 * 60_000), + createdAt: new Date(), + updatedAt: new Date(), + }, + params: {}, + }), + extendCreditLease: jest.fn().mockReturnValue(extendPending), + releaseCreditLease: jest.fn(), + }; + const { manager, store } = makeManager(creditsClient); + // Seed a live, debited lease so `maybeExtendInBackground` triggers an extend. + await store.replace({ + leaseId: "lse_live", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 1000, + expiresAt: new Date(Date.now() + 5 * 60_000), + }); + await store.tryReserve("co_1", "ct_1", 800); // 200 left → below 25% watermark + + // Kick off the extend (will hang on extendPending). + const extendP = manager.maybeExtendInBackground("co_1", "ct_1"); + + // Drop the lease so acquire is needed, then call acquireIfNeeded — + // this must NOT receive the in-flight extend promise. + await store.drop("co_1", "ct_1"); + const acquired = await manager.acquireIfNeeded("co_1", "ct_1"); + expect(acquired?.leaseId).toBe("lse_fresh"); + expect(creditsClient.acquireCreditLease).toHaveBeenCalledTimes(1); + + // Let the extend resolve so we don't leak the pending promise. + releaseExtend({ + data: { + id: "lse_live", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 2000, + expiresAt: new Date(Date.now() + 5 * 60_000), + createdAt: new Date(), + updatedAt: new Date(), + }, + params: {}, + }); + await extendP; + }); + + it("releases the redundant lease when it loses a concurrent cross-pod acquire race", async () => { + // Two managers share one backing store (mirrors two pods on one Redis). + // Both see an empty slot and acquire from the API in parallel; only one + // lease can win the slot, and the loser must release the lease it minted + // so it isn't left as an orphaned hold against the company balance. + const sharedStore: ILeaseStore = new RedisLeaseStore({ client: makeFakeRedis() }); + const config = { + defaultLeaseDuration: 5 * 60_000, + defaultReservationTTL: 60_000, + defaultLeaseSize: 1000, + lowWaterMark: 0.25, + }; + const mkManager = (leaseId: string) => { + const creditsClient = { + acquireCreditLease: jest.fn().mockResolvedValue({ + data: { + id: leaseId, + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 1000, + expiresAt: new Date(Date.now() + 5 * 60_000), + createdAt: new Date(), + updatedAt: new Date(), + }, + params: {}, + }), + extendCreditLease: jest.fn(), + releaseCreditLease: jest.fn().mockResolvedValue({ data: {}, params: {} }), + }; + const manager = new CreditLeaseManager({ + // biome-ignore lint/suspicious/noExplicitAny: stubbed client + creditsClient: creditsClient as any, + leaseStore: sharedStore, + logger: makeLogger(), + config, + }); + return { manager, creditsClient }; + }; + const a = mkManager("lse_a"); + const b = mkManager("lse_b"); + + const [ea, eb] = await Promise.all([ + a.manager.acquireIfNeeded("co_1", "ct_1"), + b.manager.acquireIfNeeded("co_1", "ct_1"), + ]); + + // Both raced to the API. + expect(a.creditsClient.acquireCreditLease).toHaveBeenCalledTimes(1); + expect(b.creditsClient.acquireCreditLease).toHaveBeenCalledTimes(1); + + // Exactly one lease survives in the shared slot, and both managers see it. + const survivor = await sharedStore.get("co_1", "ct_1"); + expect(survivor).toBeDefined(); + expect(["lse_a", "lse_b"]).toContain(survivor!.leaseId); + expect(ea?.leaseId).toBe(survivor!.leaseId); + expect(eb?.leaseId).toBe(survivor!.leaseId); + + // The loser released its redundant lease exactly once, across both. + const releaseCalls = + a.creditsClient.releaseCreditLease.mock.calls.length + + b.creditsClient.releaseCreditLease.mock.calls.length; + expect(releaseCalls).toBe(1); + const loserId = survivor!.leaseId === "lse_a" ? "lse_b" : "lse_a"; + const loser = survivor!.leaseId === "lse_a" ? b : a; + expect(loser.creditsClient.releaseCreditLease).toHaveBeenCalledWith(loserId, {}); + }); + + it("does not release the lease on a normal uncontended acquire", async () => { + const creditsClient = { + acquireCreditLease: jest.fn().mockResolvedValue({ + data: { + id: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 1000, + expiresAt: new Date(Date.now() + 5 * 60_000), + createdAt: new Date(), + updatedAt: new Date(), + }, + params: {}, + }), + extendCreditLease: jest.fn(), + releaseCreditLease: jest.fn().mockResolvedValue({ data: {}, params: {} }), + }; + const { manager } = makeManager(creditsClient); + await manager.acquireIfNeeded("co_1", "ct_1"); + expect(creditsClient.releaseCreditLease).not.toHaveBeenCalled(); + }); +}); diff --git a/tests/unit/credits/lease-store.test.ts b/tests/unit/credits/lease-store.test.ts new file mode 100644 index 00000000..97aa962e --- /dev/null +++ b/tests/unit/credits/lease-store.test.ts @@ -0,0 +1,172 @@ +import { LeaseStore } from "../../../src/credits/lease-store"; + +describe("LeaseStore", () => { + let store: LeaseStore; + + beforeEach(() => { + store = new LeaseStore(); + }); + + it("replaces installs a lease with localRemaining=grantedAmount", async () => { + await store.replace({ + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 100, + expiresAt: new Date(Date.now() + 60_000), + }); + const entry = store.get("co_1", "ct_1"); + expect(entry).toBeDefined(); + expect(entry?.localRemainingCredits).toBe(100); + expect(entry?.grantedAmount).toBe(100); + }); + + it("tryReserve debits localRemaining and returns true on success", async () => { + await store.replace({ + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 100, + expiresAt: new Date(Date.now() + 60_000), + }); + const ok = await store.tryReserve("co_1", "ct_1", 30); + expect(ok).toBe(true); + expect(store.get("co_1", "ct_1")?.localRemainingCredits).toBe(70); + }); + + it("tryReserve returns false when insufficient and does not debit", async () => { + await store.replace({ + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 100, + expiresAt: new Date(Date.now() + 60_000), + }); + const ok = await store.tryReserve("co_1", "ct_1", 150); + expect(ok).toBe(false); + expect(store.get("co_1", "ct_1")?.localRemainingCredits).toBe(100); + }); + + it("tryReserve returns false against an expired lease and does not debit", async () => { + await store.replace({ + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 100, + expiresAt: new Date(Date.now() - 1_000), + }); + const ok = await store.tryReserve("co_1", "ct_1", 10); + expect(ok).toBe(false); + // Balance untouched — an expired lease is stale, the server has released it. + expect(store.get("co_1", "ct_1")?.localRemainingCredits).toBe(100); + }); + + it("refund adds credits back, capped at grantedAmount", async () => { + await store.replace({ + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 100, + expiresAt: new Date(Date.now() + 60_000), + }); + await store.tryReserve("co_1", "ct_1", 30); + await store.refund("co_1", "ct_1", 20); + expect(store.get("co_1", "ct_1")?.localRemainingCredits).toBe(90); + // Refunding past grantedAmount caps at grantedAmount + await store.refund("co_1", "ct_1", 9999); + expect(store.get("co_1", "ct_1")?.localRemainingCredits).toBe(100); + }); + + it("concurrent tryReserves serialize and only succeed up to grantedAmount", async () => { + await store.replace({ + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 100, + expiresAt: new Date(Date.now() + 60_000), + }); + const results = await Promise.all([ + store.tryReserve("co_1", "ct_1", 40), + store.tryReserve("co_1", "ct_1", 40), + store.tryReserve("co_1", "ct_1", 40), + ]); + const successes = results.filter((r) => r === true).length; + expect(successes).toBe(2); + expect(store.get("co_1", "ct_1")?.localRemainingCredits).toBe(20); + }); + + it("extend adds to grantedAmount and localRemaining and updates expiry", async () => { + await store.replace({ + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 100, + expiresAt: new Date(Date.now() + 10_000), + }); + await store.tryReserve("co_1", "ct_1", 30); + const newExpiry = new Date(Date.now() + 60_000); + await store.extend("co_1", "ct_1", 50, newExpiry); + const e = store.get("co_1", "ct_1"); + expect(e?.grantedAmount).toBe(150); + expect(e?.localRemainingCredits).toBe(120); + expect(e?.expiresAt.getTime()).toBe(newExpiry.getTime()); + }); + + it("drop removes the lease entry", async () => { + await store.replace({ + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 100, + expiresAt: new Date(Date.now() + 60_000), + }); + await store.drop("co_1", "ct_1"); + expect(store.get("co_1", "ct_1")).toBeUndefined(); + }); + + it("replace keeps an existing live lease (preserving its debited balance) and reports it did not write", async () => { + const wroteFirst = await store.replace({ + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 100, + expiresAt: new Date(Date.now() + 60_000), + }); + expect(wroteFirst).toBe(true); + await store.tryReserve("co_1", "ct_1", 40); + + // A racing acquire installs a different lease id while the first is + // still live — the existing debited balance must win. + const wroteSecond = await store.replace({ + leaseId: "lse_2", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 100, + expiresAt: new Date(Date.now() + 60_000), + }); + expect(wroteSecond).toBe(false); + const entry = store.get("co_1", "ct_1"); + expect(entry?.leaseId).toBe("lse_1"); + expect(entry?.localRemainingCredits).toBe(60); + }); + + it("replace overwrites an expired lease and reports it wrote", async () => { + await store.replace({ + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 100, + expiresAt: new Date(Date.now() - 1_000), + }); + const wrote = await store.replace({ + leaseId: "lse_2", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 100, + expiresAt: new Date(Date.now() + 60_000), + }); + expect(wrote).toBe(true); + expect(store.get("co_1", "ct_1")?.leaseId).toBe("lse_2"); + expect(store.get("co_1", "ct_1")?.localRemainingCredits).toBe(100); + }); +}); diff --git a/tests/unit/credits/redis-lease-store.test.ts b/tests/unit/credits/redis-lease-store.test.ts new file mode 100644 index 00000000..81d6b014 --- /dev/null +++ b/tests/unit/credits/redis-lease-store.test.ts @@ -0,0 +1,161 @@ +import { RedisLeaseStore } from "../../../src/credits/redis-lease-store"; +import { makeFakeRedis } from "./fake-redis"; + +describe("RedisLeaseStore", () => { + it("replace installs a fresh lease, returns true", async () => { + const client = makeFakeRedis(); + const store = new RedisLeaseStore({ client }); + const wrote = await store.replace({ + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 100, + expiresAt: new Date(Date.now() + 60_000), + }); + expect(wrote).toBe(true); + const entry = await store.get("co_1", "ct_1"); + expect(entry?.leaseId).toBe("lse_1"); + expect(entry?.localRemainingCredits).toBe(100); + expect(entry?.grantedAmount).toBe(100); + }); + + it("replace preserves debited state when the same live leaseId is rewritten", async () => { + const client = makeFakeRedis(); + const store = new RedisLeaseStore({ client }); + await store.replace({ + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 1000, + expiresAt: new Date(Date.now() + 60_000), + }); + await store.tryReserve("co_1", "ct_1", 400); + // Simulate a second pod calling acquire and getting the same lease back. + const wrote = await store.replace({ + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 1000, + expiresAt: new Date(Date.now() + 60_000), + }); + expect(wrote).toBe(false); + const entry = await store.get("co_1", "ct_1"); + expect(entry?.localRemainingCredits).toBe(600); + }); + + it("tryReserve atomically gates against the shared balance", async () => { + const client = makeFakeRedis(); + const store = new RedisLeaseStore({ client }); + await store.replace({ + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 100, + expiresAt: new Date(Date.now() + 60_000), + }); + const results = await Promise.all([ + store.tryReserve("co_1", "ct_1", 40), + store.tryReserve("co_1", "ct_1", 40), + store.tryReserve("co_1", "ct_1", 40), + ]); + // Two should succeed, one should fail (40 left after two debits) + expect(results.filter((r) => r).length).toBe(2); + const entry = await store.get("co_1", "ct_1"); + expect(entry?.localRemainingCredits).toBe(20); + }); + + it("tryReserve rejects an expired-but-not-yet-evicted lease", async () => { + const client = makeFakeRedis(); + const store = new RedisLeaseStore({ client }); + // Expiry is in the past, but the TTL grace keeps the row alive in Redis. + await store.replace({ + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 100, + expiresAt: new Date(Date.now() - 1_000), + }); + // Row still present (within grace window)... + expect(await store.get("co_1", "ct_1")).toBeDefined(); + // ...but the expiry guard refuses to reserve against stale balance. + const ok = await store.tryReserve("co_1", "ct_1", 10); + expect(ok).toBe(false); + expect((await store.get("co_1", "ct_1"))?.localRemainingCredits).toBe(100); + }); + + it("refund caps at grantedAmount", async () => { + const client = makeFakeRedis(); + const store = new RedisLeaseStore({ client }); + await store.replace({ + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 100, + expiresAt: new Date(Date.now() + 60_000), + }); + await store.tryReserve("co_1", "ct_1", 30); + await store.refund("co_1", "ct_1", 9999); + const entry = await store.get("co_1", "ct_1"); + expect(entry?.localRemainingCredits).toBe(100); + }); + + it("replace keeps a DIFFERENT live lease id, preserving its debited balance", async () => { + // Models two pods racing the first acquire for the same slot: the loser + // must not clobber the winner's already-debited balance. + const client = makeFakeRedis(); + const store = new RedisLeaseStore({ client }); + await store.replace({ + leaseId: "lse_winner", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 1000, + expiresAt: new Date(Date.now() + 60_000), + }); + await store.tryReserve("co_1", "ct_1", 400); + const wrote = await store.replace({ + leaseId: "lse_loser", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 1000, + expiresAt: new Date(Date.now() + 60_000), + }); + expect(wrote).toBe(false); + const entry = await store.get("co_1", "ct_1"); + expect(entry?.leaseId).toBe("lse_winner"); + expect(entry?.localRemainingCredits).toBe(600); + }); + + it("extend honors defaultLeaseDurationMs option when newExpiresAt is omitted", async () => { + const client = makeFakeRedis(); + const store = new RedisLeaseStore({ client, defaultLeaseDurationMs: 250 }); + await store.replace({ + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 100, + expiresAt: new Date(Date.now() + 60_000), + }); + const before = Date.now(); + await store.extend("co_1", "ct_1", 50); + const entry = await store.get("co_1", "ct_1"); + // Configured fallback is 250ms — expiry should land near `before + 250`. + const expiry = entry?.expiresAt.getTime() ?? 0; + expect(expiry).toBeGreaterThanOrEqual(before + 200); + expect(expiry).toBeLessThanOrEqual(before + 500); + expect(entry?.grantedAmount).toBe(150); + }); + + it("drop removes the lease hash", async () => { + const client = makeFakeRedis(); + const store = new RedisLeaseStore({ client }); + await store.replace({ + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 100, + expiresAt: new Date(Date.now() + 60_000), + }); + await store.drop("co_1", "ct_1"); + expect(await store.get("co_1", "ct_1")).toBeUndefined(); + }); +}); diff --git a/tests/unit/credits/redis-reservation-store.test.ts b/tests/unit/credits/redis-reservation-store.test.ts new file mode 100644 index 00000000..f24215af --- /dev/null +++ b/tests/unit/credits/redis-reservation-store.test.ts @@ -0,0 +1,246 @@ +import { RedisLeaseStore } from "../../../src/credits/redis-lease-store"; +import { RedisReservationStore } from "../../../src/credits/redis-reservation-store"; +import type { Reservation } from "../../../src/credits/types"; +import { makeFakeRedis } from "./fake-redis"; + +function makeReservation(overrides: Partial = {}): Reservation { + return { + id: "res_1", + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + eventSubtype: "inference_tokens", + quantityReserved: 10, + creditsReserved: 100, + consumptionRate: 10, + expiresAt: new Date(Date.now() + 60_000), + evalCtx: { company: { id: "co_1" } }, + ...overrides, + }; +} + +describe("RedisReservationStore", () => { + it("add round-trips through get and shows up in the expiry index", async () => { + const client = makeFakeRedis(); + const leaseStore = new RedisLeaseStore({ client }); + const reservations = new RedisReservationStore({ client, leaseStore, sweepIntervalMs: 60_000 }); + await leaseStore.replace({ + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 1000, + expiresAt: new Date(Date.now() + 60_000), + }); + + const reservation = makeReservation(); + await reservations.add(reservation); + const fetched = await reservations.get(reservation.id); + expect(fetched?.creditsReserved).toBe(100); + expect(await reservations.size()).toBe(1); + reservations.stop(); + }); + + it("consume refunds unused credits atomically", async () => { + const client = makeFakeRedis(); + const leaseStore = new RedisLeaseStore({ client }); + const reservations = new RedisReservationStore({ client, leaseStore, sweepIntervalMs: 60_000 }); + await leaseStore.replace({ + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 1000, + expiresAt: new Date(Date.now() + 60_000), + }); + await leaseStore.tryReserve("co_1", "ct_1", 100); + await reservations.add(makeReservation()); + expect((await leaseStore.get("co_1", "ct_1"))?.localRemainingCredits).toBe(900); + + const consumed = await reservations.consume("res_1", 30); + expect(consumed).toBe(30); + // Refunded 100-30=70 + expect((await leaseStore.get("co_1", "ct_1"))?.localRemainingCredits).toBe(970); + expect(await reservations.get("res_1")).toBeUndefined(); + reservations.stop(); + }); + + it("sweepExpired returns expired reservations to the lease", async () => { + const client = makeFakeRedis(); + const leaseStore = new RedisLeaseStore({ client }); + const reservations = new RedisReservationStore({ client, leaseStore, sweepIntervalMs: 60_000 }); + await leaseStore.replace({ + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 1000, + expiresAt: new Date(Date.now() + 60_000), + }); + await leaseStore.tryReserve("co_1", "ct_1", 100); + await reservations.add(makeReservation({ expiresAt: new Date(Date.now() - 1) })); + + const swept = await reservations.sweepExpired(); + expect(swept).toBe(1); + expect((await leaseStore.get("co_1", "ct_1"))?.localRemainingCredits).toBe(1000); + expect(await reservations.get("res_1")).toBeUndefined(); + reservations.stop(); + }); + + it("reservedCredits sums open reservations and drops consumed ones", async () => { + const client = makeFakeRedis(); + const leaseStore = new RedisLeaseStore({ client }); + const reservations = new RedisReservationStore({ client, leaseStore, sweepIntervalMs: 60_000 }); + await leaseStore.replace({ + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 1000, + expiresAt: new Date(Date.now() + 60_000), + }); + await leaseStore.tryReserve("co_1", "ct_1", 350); + await reservations.add(makeReservation({ id: "res_1", creditsReserved: 100 })); + await reservations.add(makeReservation({ id: "res_2", creditsReserved: 250 })); + // Different credit shouldn't count toward ct_1. + await reservations.add(makeReservation({ id: "res_3", creditTypeId: "ct_2", creditsReserved: 999 })); + + expect(await reservations.reservedCredits("co_1", "ct_1")).toBe(350); + expect(await reservations.reservedCredits("co_1", "ct_2")).toBe(999); + + await reservations.consume("res_1", 40); + expect(await reservations.reservedCredits("co_1", "ct_1")).toBe(250); + reservations.stop(); + }); + + it("never issues a multi-key Lua EVAL (Cluster-safe: no CROSSSLOT)", async () => { + // A multi-key script whose keys hash to different slots raises CROSSSLOT + // on Redis Cluster. Guard that every EVAL the lease + reservation stores + // send touches exactly one key, across the full lifecycle. + const base = makeFakeRedis(); + const evalKeyCounts: number[] = []; + const client = { + ...base, + eval: (script: string, options: { keys: string[]; arguments: string[] }) => { + evalKeyCounts.push(options.keys.length); + return base.eval(script, options); + }, + } as typeof base; + + const leaseStore = new RedisLeaseStore({ client }); + const reservations = new RedisReservationStore({ client, leaseStore, sweepIntervalMs: 60_000 }); + + // Exercise every Lua path: replace, tryReserve, refund, extend (lease) + // and add, consume, sweep (reservation, whose refund delegates to the + // lease store's single-key REFUND). + await leaseStore.replace({ + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 1000, + expiresAt: new Date(Date.now() + 60_000), + }); + await leaseStore.tryReserve("co_1", "ct_1", 200); + await leaseStore.refund("co_1", "ct_1", 50); + await leaseStore.extend("co_1", "ct_1", 100, new Date(Date.now() + 120_000)); + await reservations.add(makeReservation({ id: "res_a", creditsReserved: 100 })); + await reservations.add(makeReservation({ id: "res_b", creditsReserved: 80, expiresAt: new Date(Date.now() - 1) })); + await reservations.consume("res_a", 40); + await reservations.sweepExpired(); + + expect(evalKeyCounts.length).toBeGreaterThan(0); + expect(evalKeyCounts.every((n) => n === 1)).toBe(true); + reservations.stop(); + }); + + it("sweep reconciles orphaned indexes when the reservation hash TTL-evicts before sweeping", async () => { + // Reproduces the failure mode where a sweeper goes silent (deploy/restart, + // event-loop starvation) long enough for Redis to evict the reservation + // hash via its TTL grace. Before the fix, the sweeper's `consume` would + // hit the hash-gone path and never clean the byExpiry set or the per-tenant + // `byCredit` hash — orphaning the index entry (unbounded zset growth) and + // permanently inflating `reservedCredits`. + jest.useFakeTimers(); + try { + const t0 = new Date("2026-01-01T00:00:00Z").getTime(); + jest.setSystemTime(t0); + + const client = makeFakeRedis(); + const leaseStore = new RedisLeaseStore({ client }); + const reservations = new RedisReservationStore({ client, leaseStore, sweepIntervalMs: 60_000 }); + await leaseStore.replace({ + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 1000, + expiresAt: new Date(t0 + 10 * 60_000), + }); + await leaseStore.tryReserve("co_1", "ct_1", 100); + // Reservation expires in 1s; its hash is reaped 30s (RES_TTL_GRACE_MS) + // later. The lease lives well past both. + await reservations.add(makeReservation({ expiresAt: new Date(t0 + 1000) })); + expect(await reservations.size()).toBe(1); + expect(await reservations.reservedCredits("co_1", "ct_1")).toBe(100); + + // Jump past the hash's TTL grace (but not the lease) so Redis would + // have evicted only the reservation hash. + jest.setSystemTime(t0 + 1000 + 30_000 + 1); + + const swept = await reservations.sweepExpired(); + // Nothing is "swept" in the refund sense — the hash is already gone, so + // the unspent slice is left for server-side lease expiry to reclaim. + expect(swept).toBe(0); + // But the orphaned secondary indexes are reconciled: no tombstone left + // in the expiry set, and reservedCredits no longer counts the evicted + // hold (this is what regressed before the fix). + expect(await reservations.size()).toBe(0); + expect(await reservations.reservedCredits("co_1", "ct_1")).toBe(0); + reservations.stop(); + } finally { + jest.useRealTimers(); + } + }); + + it("sweep handles a legacy bare-id index member (rolling upgrade)", async () => { + // A pre-upgrade SDK indexed reservations under the bare id, not the + // `company|credit|id` tuple. The sweeper must still consume + refund such + // entries and remove them from the set. + const client = makeFakeRedis(); + const leaseStore = new RedisLeaseStore({ client }); + const reservations = new RedisReservationStore({ client, leaseStore, sweepIntervalMs: 60_000 }); + await leaseStore.replace({ + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 1000, + expiresAt: new Date(Date.now() + 60_000), + }); + await leaseStore.tryReserve("co_1", "ct_1", 100); + await reservations.add(makeReservation({ expiresAt: new Date(Date.now() - 1) })); + // Rewrite the index entry to the legacy bare-id form. + await client.zRem("schematic:credit-reservations:byExpiry", "co_1|ct_1|res_1"); + await client.zAdd("schematic:credit-reservations:byExpiry", { score: Date.now() - 1, value: "res_1" }); + + const swept = await reservations.sweepExpired(); + expect(swept).toBe(1); + expect((await leaseStore.get("co_1", "ct_1"))?.localRemainingCredits).toBe(1000); + expect(await reservations.size()).toBe(0); + expect(await reservations.reservedCredits("co_1", "ct_1")).toBe(0); + reservations.stop(); + }); + + it("double-consume returns null on the second call", async () => { + const client = makeFakeRedis(); + const leaseStore = new RedisLeaseStore({ client }); + const reservations = new RedisReservationStore({ client, leaseStore, sweepIntervalMs: 60_000 }); + await leaseStore.replace({ + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 1000, + expiresAt: new Date(Date.now() + 60_000), + }); + await leaseStore.tryReserve("co_1", "ct_1", 100); + await reservations.add(makeReservation()); + await reservations.consume("res_1", 50); + const second = await reservations.consume("res_1", 10); + expect(second).toBeNull(); + reservations.stop(); + }); +}); diff --git a/tests/unit/credits/reservation-store.test.ts b/tests/unit/credits/reservation-store.test.ts new file mode 100644 index 00000000..ab793026 --- /dev/null +++ b/tests/unit/credits/reservation-store.test.ts @@ -0,0 +1,107 @@ +import { LeaseStore } from "../../../src/credits/lease-store"; +import { ReservationStore } from "../../../src/credits/reservation-store"; +import type { Reservation } from "../../../src/credits/types"; + +function makeReservation(overrides: Partial = {}): Reservation { + return { + id: "res_1", + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + eventSubtype: "inference_tokens", + quantityReserved: 10, + creditsReserved: 100, + consumptionRate: 10, + expiresAt: new Date(Date.now() + 60_000), + evalCtx: { company: { id: "co_1" } }, + ...overrides, + }; +} + +describe("ReservationStore", () => { + let leases: LeaseStore; + let reservations: ReservationStore; + + beforeEach(async () => { + leases = new LeaseStore(); + reservations = new ReservationStore(leases, 50); + await leases.replace({ + leaseId: "lse_1", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 1000, + expiresAt: new Date(Date.now() + 60_000), + }); + }); + + afterEach(() => { + reservations.stop(); + }); + + it("consume refunds unused credits to the lease", async () => { + await leases.tryReserve("co_1", "ct_1", 100); + const reservation = makeReservation(); + reservations.add(reservation); + expect(leases.get("co_1", "ct_1")?.localRemainingCredits).toBe(900); + + const consumed = await reservations.consume(reservation.id, 30); + expect(consumed).toBe(30); + // Refunded 100 - 30 = 70 + expect(leases.get("co_1", "ct_1")?.localRemainingCredits).toBe(970); + expect(reservations.get(reservation.id)).toBeUndefined(); + }); + + it("consume clamps credits consumed to creditsReserved", async () => { + await leases.tryReserve("co_1", "ct_1", 100); + const reservation = makeReservation(); + reservations.add(reservation); + const consumed = await reservations.consume(reservation.id, 999); + expect(consumed).toBe(100); + expect(leases.get("co_1", "ct_1")?.localRemainingCredits).toBe(900); + }); + + it("consume returns null on missing reservation", async () => { + const result = await reservations.consume("nope", 10); + expect(result).toBeNull(); + }); + + it("sweepExpired refunds expired reservations to the lease", async () => { + await leases.tryReserve("co_1", "ct_1", 100); + const reservation = makeReservation({ expiresAt: new Date(Date.now() - 1) }); + reservations.add(reservation); + + const swept = await reservations.sweepExpired(); + expect(swept).toBe(1); + expect(reservations.get(reservation.id)).toBeUndefined(); + expect(leases.get("co_1", "ct_1")?.localRemainingCredits).toBe(1000); + }); + + it("sweepExpired ignores non-expired reservations", async () => { + await leases.tryReserve("co_1", "ct_1", 100); + reservations.add(makeReservation()); + const swept = await reservations.sweepExpired(); + expect(swept).toBe(0); + expect(reservations.size()).toBe(1); + }); + + it("reservedCredits sums open reservations for a (company, credit)", async () => { + reservations.add(makeReservation({ id: "res_1", creditsReserved: 100 })); + reservations.add(makeReservation({ id: "res_2", creditsReserved: 250 })); + // Different credit / company shouldn't count. + reservations.add(makeReservation({ id: "res_3", creditTypeId: "ct_2", creditsReserved: 999 })); + reservations.add(makeReservation({ id: "res_4", companyId: "co_2", creditsReserved: 999 })); + + expect(reservations.reservedCredits("co_1", "ct_1")).toBe(350); + expect(reservations.reservedCredits("co_1", "ct_2")).toBe(999); + expect(reservations.reservedCredits("co_unknown", "ct_1")).toBe(0); + }); + + it("reservedCredits drops a reservation once it is consumed", async () => { + await leases.tryReserve("co_1", "ct_1", 100); + reservations.add(makeReservation()); + expect(reservations.reservedCredits("co_1", "ct_1")).toBe(100); + + await reservations.consume("res_1", 30); + expect(reservations.reservedCredits("co_1", "ct_1")).toBe(0); + }); +}); diff --git a/tests/unit/credits/wasm-credit-gate.test.ts b/tests/unit/credits/wasm-credit-gate.test.ts new file mode 100644 index 00000000..86b4aef7 --- /dev/null +++ b/tests/unit/credits/wasm-credit-gate.test.ts @@ -0,0 +1,227 @@ +/** + * Real-WASM tests for the credit-lease check path. + * + * Everything else in `tests/unit/credits` stubs the rules engine, so the + * contract that actually matters most — substituting the lease balance into + * the company's `creditBalances` and letting the WASM's `event_usage` gate + * decide allow/deny — is never exercised against the real engine. These tests + * load the bundled WASM (`src/wasm/rulesengine.js`) and drive a credit-balance + * flag end-to-end through `checkWithLease`, plus a direct rules-engine check, so + * that a drift in the snake_case option envelope (`event_usage`) or in the + * camelCase entity shape the SDK now feeds the engine would fail here instead of + * silently mis-gating in production. + */ +import type * as api from "../../../src/api"; +import type { DataStreamClient } from "../../../src/datastream"; +import type { Logger } from "../../../src/logger"; +import { RulesEngineClient } from "../../../src/rules-engine"; +import { LeaseStore } from "../../../src/credits/lease-store"; +import { ReservationStore } from "../../../src/credits/reservation-store"; +import { CreditLeaseManager } from "../../../src/credits/lease-manager"; +import { checkWithLease, type CreditCheckDeps } from "../../../src/credits/check"; +import type { CheckResult } from "../../../src/credits/types"; + +const silentLogger: Logger = { + error: () => {}, + warn: () => {}, + info: () => {}, + debug: () => {}, +}; + +const CREDIT_ID = "credit-1"; +const EVENT_SUBTYPE = "inference_tokens"; + +// Credit condition (camelCase — the shape the datastream cache now holds after +// the serializer canonicalizes wire payloads). `consumptionRate` 1 keeps +// credits == quantity so the math is easy to follow. +const creditCondition = { + id: "cond-credit", + accountId: "acct", + environmentId: "env", + conditionType: "credit", + operator: "gte", + resourceIds: [] as string[], + traitValue: "", + metricValue: 0, + creditId: CREDIT_ID, + consumptionRate: 1, + eventSubtype: EVENT_SUBTYPE, +}; + +// A `company`-membership condition we can flip to force the WASM to deny the +// rule for reasons unrelated to the credit balance (so the lease reservation +// gets refunded). +function companyCondition(resourceIds: string[]) { + return { + id: "cond-company", + accountId: "acct", + environmentId: "env", + conditionType: "company", + operator: "eq", + resourceIds, + traitValue: "", + metricValue: 0, + }; +} + +function creditFlag(extraConditions: object[] = []): api.RulesengineFlag { + return { + id: "flag-infer", + accountId: "acct", + environmentId: "env", + key: "infer", + defaultValue: false, + rules: [ + { + id: "rule-credit", + accountId: "acct", + environmentId: "env", + name: "Credit", + ruleType: "plan_entitlement", + priority: 100, + value: true, + conditions: [creditCondition, ...extraConditions], + conditionGroups: [], + }, + ], + } as unknown as api.RulesengineFlag; +} + +function company(creditBalance: number): api.RulesengineCompany { + return { + id: "co", + accountId: "acct", + environmentId: "env", + keys: {}, + creditBalances: { [CREDIT_ID]: creditBalance }, + } as unknown as api.RulesengineCompany; +} + +// Minimal datastream stub: serves a fixed flag + company from "cache" and hands +// back the REAL rules engine. +function fakeDatastream( + engine: RulesEngineClient, + flag: api.RulesengineFlag, + co: api.RulesengineCompany, +): DataStreamClient { + return { + getFlag: async () => flag, + getCachedCompany: async () => co, + getCachedUser: async () => null, + getRulesEngine: () => engine, + } as unknown as DataStreamClient; +} + +function makeDeps( + engine: RulesEngineClient, + flag: api.RulesengineFlag, + co: api.RulesengineCompany, + grantedAmount: number, +): { deps: CreditCheckDeps; leaseStore: LeaseStore; reservations: ReservationStore } { + const leaseStore = new LeaseStore(); + const reservations = new ReservationStore(leaseStore, 60_000); + const creditsClient = { + acquireCreditLease: jest.fn().mockResolvedValue({ + data: { + id: "lse-1", + companyId: "co", + creditTypeId: CREDIT_ID, + grantedAmount, + expiresAt: new Date(Date.now() + 60_000), + }, + }), + extendCreditLease: jest.fn(), + releaseCreditLease: jest.fn().mockResolvedValue({}), + // eslint-disable-next-line @typescript-eslint/no-explicit-any + } as any; + const manager = new CreditLeaseManager({ + creditsClient, + leaseStore, + logger: silentLogger, + config: { defaultLeaseSize: grantedAmount, defaultLeaseDuration: 60_000 }, + }); + const deps: CreditCheckDeps = { + leaseStore, + reservations, + manager, + datastream: fakeDatastream(engine, flag, co), + logger: silentLogger, + }; + return { deps, leaseStore, reservations }; +} + +const evalCtx: api.CheckFlagRequestBody = { company: { id: "co" } }; +const failFallback = (): Promise => { + throw new Error("fallback should not be reached on the lease path"); +}; + +describe("credit-lease check against the real WASM engine", () => { + let engine: RulesEngineClient; + + beforeAll(async () => { + engine = new RulesEngineClient(); + await engine.initialize(); + }); + + it("substitutes the lease balance and lets a within-balance usage through, issuing a reservation", async () => { + const { deps, leaseStore, reservations } = makeDeps(engine, creditFlag(), company(100), 10_000); + + const result = await checkWithLease( + deps, + "infer", + evalCtx, + { eventUsage: { [EVENT_SUBTYPE]: 50 } }, + failFallback, + ); + + expect(result.allowed).toBe(true); + expect(result.value).toBe(true); + expect(result.reservation).toBeDefined(); + expect(result.reservation?.creditTypeId).toBe(CREDIT_ID); + expect(result.reservation?.creditsReserved).toBe(50); + // The 50-credit reservation stays debited from the lease's local view. + expect((await leaseStore.get("co", CREDIT_ID))?.localRemainingCredits).toBe(9_950); + expect(reservations.size()).toBe(1); + }); + + it("denies (and refunds the reservation) when the WASM rule fails for a non-credit reason", async () => { + // Credit balance is plentiful, but the company-membership condition + // excludes this company, so the WASM denies the rule. The reservation + // taken before the eval must be returned to the lease. + const flag = creditFlag([companyCondition(["some-other-company"])]); + const { deps, leaseStore, reservations } = makeDeps(engine, flag, company(10_000), 10_000); + + const result = await checkWithLease( + deps, + "infer", + evalCtx, + { eventUsage: { [EVENT_SUBTYPE]: 50 } }, + failFallback, + ); + + expect(result.allowed).toBe(false); + expect(result.value).toBe(false); + expect(result.reservation).toBeUndefined(); + // Reservation refunded: lease back to full, table empty. + expect((await leaseStore.get("co", CREDIT_ID))?.localRemainingCredits).toBe(10_000); + expect(reservations.size()).toBe(0); + }); + + it("serializes CheckFlagOptions to the snake_case event_usage envelope the WASM gates on", async () => { + // Direct rules-engine check: prove the SDK's camelCase `eventUsage` + // option reaches the WASM as `event_usage` and gates exactly at the + // balance boundary. This is the contract `checkWithLease` leans on. + const flag = creditFlag(); + const co = company(100); + + const under = await engine.checkFlagWithOptions(flag, co, null, { + eventUsage: { [EVENT_SUBTYPE]: 50 }, + }); + const over = await engine.checkFlagWithOptions(flag, co, null, { + eventUsage: { [EVENT_SUBTYPE]: 150 }, + }); + + expect(under.value).toBe(true); + expect(over.value).toBe(false); + }); +}); diff --git a/tests/unit/datastream/datastream-client.test.ts b/tests/unit/datastream/datastream-client.test.ts index da27a0c0..984096e2 100644 --- a/tests/unit/datastream/datastream-client.test.ts +++ b/tests/unit/datastream/datastream-client.test.ts @@ -40,11 +40,23 @@ jest.mock('../../../src/datastream/websocket-client', () => { }; }); -// Mock RulesEngineClient so we can control what checkFlag returns -const mockRulesEngineInstance = { +// Mock RulesEngineClient so we can control what checkFlag returns. +// The real client routes evaluateFlag through `checkFlagWithOptions`; mock +// it to delegate to `checkFlag` so existing tests that stub `checkFlag` keep +// working without per-test churn. +const mockRulesEngineInstance: { + initialize: jest.Mock; + isInitialized: jest.Mock; + checkFlag: jest.Mock; + checkFlagWithOptions: jest.Mock; + getVersionKey: jest.Mock; +} = { initialize: jest.fn().mockResolvedValue(undefined), isInitialized: jest.fn().mockReturnValue(false), checkFlag: jest.fn(), + checkFlagWithOptions: jest.fn((flag, company, user, _options) => + mockRulesEngineInstance.checkFlag(flag, company, user), + ), getVersionKey: jest.fn().mockReturnValue('1'), }; diff --git a/tests/unit/wrapper.test.ts b/tests/unit/wrapper.test.ts index e7033088..6d961a1f 100644 --- a/tests/unit/wrapper.test.ts +++ b/tests/unit/wrapper.test.ts @@ -23,6 +23,7 @@ jest.mock("../../src/events", () => { return { EventBuffer: jest.fn().mockImplementation(() => ({ push: jest.fn(), + flush: jest.fn().mockResolvedValue(undefined), stop: jest.fn().mockResolvedValue(undefined), })), }; @@ -188,7 +189,6 @@ describe("SchematicClient wrapper - flag checking behavior", () => { await client.close(); }); }); - describe("event options", () => { // eslint-disable-next-line @typescript-eslint/no-require-imports const { EventBuffer } = require("../../src/events"); @@ -256,6 +256,166 @@ describe("SchematicClient wrapper - flag checking behavior", () => { await client.close(); }); }); + + describe("identify with prewarm", () => { + it("forwards prewarm credit type ids to client.prewarm and flushes the buffer", async () => { + const client = new SchematicClient({ + apiKey: "test-api-key", + cacheProviders: { flagChecks: [] }, + logger: mockLogger, + }); + const prewarmSpy = jest + .spyOn(client, "prewarm") + .mockResolvedValue(undefined); + // Reach into the buffer mock to verify flush is triggered so the + // server picks up the identify event before prewarm starts polling. + // biome-ignore lint/suspicious/noExplicitAny: introspect mock + const flushMock = (client as any).eventBuffer.flush as jest.Mock; + + await client.identify( + { + keys: { id: "user-1" }, + company: { keys: { id: "comp-1" } }, + }, + { prewarm: ["credit-type-1", "credit-type-2"] }, + ); + + // Yield once so the fire-and-forget prewarm resolves. + await new Promise((r) => setImmediate(r)); + + expect(flushMock).toHaveBeenCalledTimes(1); + expect(prewarmSpy).toHaveBeenCalledWith( + { company: { id: "comp-1" }, user: { id: "user-1" } }, + ["credit-type-1", "credit-type-2"], + ); + + await client.close(); + }); + + it("does not call prewarm or flush when options.prewarm is omitted", async () => { + const client = new SchematicClient({ + apiKey: "test-api-key", + cacheProviders: { flagChecks: [] }, + logger: mockLogger, + }); + const prewarmSpy = jest + .spyOn(client, "prewarm") + .mockResolvedValue(undefined); + // biome-ignore lint/suspicious/noExplicitAny: introspect mock + const flushMock = (client as any).eventBuffer.flush as jest.Mock; + + await client.identify({ + keys: { id: "user-1" }, + company: { keys: { id: "comp-1" } }, + }); + + await new Promise((r) => setImmediate(r)); + expect(prewarmSpy).not.toHaveBeenCalled(); + expect(flushMock).not.toHaveBeenCalled(); + + await client.close(); + }); + }); +}); + +describe("SchematicClient wrapper - getCreditBalance lease awareness", () => { + const mockLogger = { + error: jest.fn(), + warn: jest.fn(), + info: jest.fn(), + debug: jest.fn(), + }; + + const companyKeys = { id: "comp-1" }; + const creditId = "credit-ai"; + + // Builds a client wired with a cached snapshot carrying `serverBalance`, a + // lease store returning `leaseEntry` (or undefined for "no lease"), and a + // reservation store summing to `reservedCredits`. All are private fields + // the wrapper reads in getCreditBalance. + const buildClient = ( + serverBalance: number, + leaseEntry: { localRemainingCredits: number; expiresAt: Date } | undefined, + reservedCredits = 0, + ) => { + const client = new SchematicClient({ offline: true, logger: mockLogger }); + (client as any).datastreamClient = { + getCachedCompany: jest.fn().mockResolvedValue({ + id: "comp-internal-id", + creditBalances: { [creditId]: serverBalance }, + }), + close: jest.fn(), + }; + (client as any).leaseStore = { + get: jest.fn().mockResolvedValue(leaseEntry), + }; + (client as any).reservations = { + reservedCredits: jest.fn().mockReturnValue(reservedCredits), + stop: jest.fn(), + }; + return client; + }; + + afterEach(() => { + jest.clearAllMocks(); + }); + + it("reports a live lease's unspent local balance as reserved on top of remaining", async () => { + // remaining (B−L)=600, lease hold (L−spent)=900 ⇒ settled 1500. + const client = buildClient(600, { + localRemainingCredits: 900, + expiresAt: new Date(Date.now() + 60_000), + }); + + const result = await client.getCreditBalance(companyKeys, creditId); + + expect(result).toEqual({ remaining: 600, reserved: 900, settled: 1500 }); + await client.close(); + }); + + it("folds open reservations into reserved so an in-flight hold doesn't dip settled", async () => { + // A 1500-credit reservation has carved out of the lease's local + // balance (900 = 2400 granted − 1500 reserved). `reserved` reports the + // whole unsettled lease hold (lease remainder 900 + open reservation + // 1500 = 2400), so settled = 600 + 2400 = 3000 = B − spent (nothing + // settled yet) and stays put whether or not the hold is open. + const client = buildClient( + 600, + { localRemainingCredits: 900, expiresAt: new Date(Date.now() + 60_000) }, + 1500, + ); + + const result = await client.getCreditBalance(companyKeys, creditId); + + expect(result).toEqual({ remaining: 600, reserved: 2400, settled: 3000 }); + await client.close(); + }); + + it("ignores an expired lease so the refunded hold isn't double-counted", async () => { + // Past expiresAt the server has swept and refunded the hold, so the + // whole balance already sits in `remaining`. Counting the dead hold too + // would report settled 23000 — the stale-overstated-badge bug. Any open + // reservation is ignored for the same reason. + const client = buildClient( + 12_000, + { localRemainingCredits: 11_000, expiresAt: new Date(Date.now() - 1_000) }, + 500, + ); + + const result = await client.getCreditBalance(companyKeys, creditId); + + expect(result).toEqual({ remaining: 12_000, reserved: 0, settled: 12_000 }); + await client.close(); + }); + + it("returns the server balance as remaining when no lease is held", async () => { + const client = buildClient(450, undefined); + + const result = await client.getCreditBalance(companyKeys, creditId); + + expect(result).toEqual({ remaining: 450, reserved: 0, settled: 450 }); + await client.close(); + }); }); describe("SchematicClient wrapper - logger configuration", () => { @@ -321,4 +481,80 @@ describe("SchematicClient wrapper - logger configuration", () => { await client.close(); }); -}); \ No newline at end of file +}); +describe("SchematicClient wrapper - credit lease store backend selection", () => { + const mockLogger = { + error: jest.fn(), + warn: jest.fn(), + info: jest.fn(), + debug: jest.fn(), + }; + + afterEach(() => { + jest.clearAllMocks(); + }); + + it("reuses the DataStream Redis client for lease state when no creditLeases.redisClient is set", async () => { + const { makeFakeRedis } = await import("./credits/fake-redis"); + const redisClient = makeFakeRedis(); + const client = new SchematicClient({ + apiKey: "test-key", + logger: mockLogger, + creditLeases: {}, + dataStream: { redisClient }, + }); + // The shared Redis backend must back leases automatically — no second + // client to wire up — so both stores are the Redis-backed variants. + expect((client as any).leaseStore?.constructor?.name).toBe("RedisLeaseStore"); + expect((client as any).reservations?.constructor?.name).toBe("RedisReservationStore"); + // No degrade warning when a shared backend is present. + expect(mockLogger.warn).not.toHaveBeenCalledWith( + expect.stringContaining("creditLeases is enabled without a shared Redis backend"), + ); + (client as any).reservations?.stop?.(); + }); + + it("falls back to in-memory stores and warns when no Redis backend is configured", async () => { + const client = new SchematicClient({ + apiKey: "test-key", + logger: mockLogger, + creditLeases: {}, + }); + expect((client as any).leaseStore?.constructor?.name).toBe("LeaseStore"); + expect((client as any).reservations?.constructor?.name).toBe("ReservationStore"); + expect(mockLogger.warn).toHaveBeenCalledWith( + expect.stringContaining("creditLeases is enabled without a shared Redis backend"), + ); + (client as any).reservations?.stop?.(); + }); + + it("close() does not release outstanding leases (they reclaim via expiry, not shutdown)", async () => { + const { makeFakeRedis } = await import("./credits/fake-redis"); + const redisClient = makeFakeRedis(); + const client = new SchematicClient({ + apiKey: "test-key", + logger: mockLogger, + creditLeases: {}, + dataStream: { redisClient }, + }); + // A shared lease lives in the backend (could have been installed by this + // pod or a sibling). Releasing it on this pod's shutdown would pull the + // grant out from under siblings — so close() must leave it alone. + const leaseStore = (client as any).leaseStore; + await leaseStore.replace({ + leaseId: "lse_shared", + companyId: "co_1", + creditTypeId: "ct_1", + grantedAmount: 1000, + expiresAt: new Date(Date.now() + 5 * 60_000), + }); + + await client.close(); + + // The old close() released + dropped every lease in the shared backend; + // the lease must now survive shutdown so siblings keep drawing on it. + const survivor = await leaseStore.get("co_1", "ct_1"); + expect(survivor).toBeDefined(); + expect(survivor?.leaseId).toBe("lse_shared"); + }); +});