From 2fef7e69810f3af3476a0aec3d88676e0a9e6247 Mon Sep 17 00:00:00 2001 From: Ben Papillon Date: Sat, 20 Jun 2026 19:18:24 -0700 Subject: [PATCH 1/2] adopt entitlement-first credit resolution + multi-unit preflight MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the structural credit-condition scan and the reserve-then-cancel round-trip in the lease path with a single entitlement-first probe: read creditId / consumptionRate / eventSubtype straight off the matched plan entitlement (immune to flag-shape drift), and skip lease + reserve entirely when the match isn't credit-metered. Gate the substituted eval with credit_cost (quantity x consumptionRate) for multi-unit preflight. Remove the bespoke getCreditBalance method — the lease-aware remaining/reserved/settled split now rides on the entitlement from a check. --- src/credits/check.ts | 238 ++++-------- src/credits/lease-manager.ts | 7 +- src/credits/types.ts | 20 +- src/rules-engine.ts | 4 + src/wrapper.ts | 75 ---- tests/unit/credits/check-and-track.test.ts | 394 +++++++++----------- tests/unit/credits/wasm-credit-gate.test.ts | 71 ++-- tests/unit/wrapper.test.ts | 96 ----- 8 files changed, 317 insertions(+), 588 deletions(-) diff --git a/src/credits/check.ts b/src/credits/check.ts index a18a1ab9..dd0dc331 100644 --- a/src/credits/check.ts +++ b/src/credits/check.ts @@ -1,10 +1,11 @@ import { randomUUID } from "crypto"; import type * as api from "../api"; +import { RulesengineEntitlementValueType } from "../api"; import type { CreditsClient } from "../api/resources/credits/client/Client"; import type { DataStreamClient } from "../datastream"; import type { Logger } from "../logger"; -import type { WasmFeatureEntitlement } from "../rules-engine"; +import type { WasmCheckFlagResult, WasmFeatureEntitlement } from "../rules-engine"; import type { CheckFlagOptions } from "../wrapper"; import { CreditLeaseManager } from "./lease-manager"; @@ -55,15 +56,19 @@ function emitFlagCheck( * Drives a single `client.check` with `usage` set. * * Steps: - * 1. Resolve the matching credit-balance condition on the flag to get - * `creditId` + `consumptionRate`. + * 1. Probe the engine once (real balance, no preflight, no substitution) and + * read the matched plan entitlement. If it isn't credit-metered + * (`valueType != credit` — boolean/override grant, numeric allocation, + * unlimited, or not entitled) there's nothing to lease: defer to the plain + * check. Otherwise `creditId` + `consumptionRate` + `eventSubtype` come + * straight off the entitlement — no structural flag scan. * 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. + * reservation we just made was debited), gating with `credit_cost` 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, @@ -151,35 +156,62 @@ export async function checkWithLease( } } - // 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`, - ); + // Entitlement-first resolution. One engine probe against the company's real + // balance — no preflight, no balance substitution — surfaces the matched + // plan entitlement. We read its *shape* only: `valueType` says whether a + // credit is metered at all, and a credit entitlement carries `creditId` + + // `consumptionRate` + `eventSubtype` directly. This replaces the structural + // credit-condition scan (immune to flag-shape drift) and lets a non-credit + // grant skip the lease + reserve round-trip entirely. + // + // The probe deliberately omits preflight: applying a credit cost to the + // (lease-depleted) server balance could fail the credit condition, drop the + // engine to a lower-priority rule, and hide the very entitlement we need to + // identify. The real preflight-aware gate runs in step 4 against the + // substituted lease balance. + let probe: WasmCheckFlagResult; + try { + probe = await datastream.getRulesEngine().checkFlagWithOptions(flag, company, user ?? null, null); + } catch (err) { + // The probe is a resolution step, not the gate — like a flag/company/user + // miss, a failure here means we couldn't resolve the credit, so defer to + // the plain check (which has its own defaultValue degradation) rather than + // hard-denying. No reservation exists yet, so nothing to cancel. + logger.warn(`Lease check: entitlement probe failed for flag ${key} (${err}), 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`); + const entitlement = probe.entitlement; + + // Not credit-metered — boolean/override grant, numeric allocation, + // unlimited, or simply not entitled. The feature resolves without drawing a + // credit, so skip the lease + reserve round-trip and let the plain check + // (preflight-aware, and the source of this check's flag_check event) decide. + // This is the path that previously cost an override-granted company a + // reserve-then-cancel. + if (!entitlement || entitlement.valueType !== RulesengineEntitlementValueType.Credit) { + logger.debug( + `Lease check: flag ${key} matched a non-credit entitlement (valueType=${ + entitlement?.valueType ?? "" + }), falling back to plain check — no reservation`, + ); return fallback(); } - // The reservation settles into a Track event named by this subtype; an - // empty event name would consume the reservation while billing nothing, - // so a condition with no resolvable subtype is treated like one with no - // consumption_rate. - const resolvedSubtype = eventSubtype ?? match.condition.eventSubtype; - if (!resolvedSubtype) { - logger.debug(`Lease check: credit condition has no event subtype, falling back`); + + const creditId = entitlement.creditId; + const consumptionRate = entitlement.consumptionRate ?? 0; + // The caller's explicit subtype wins; otherwise the entitlement names the + // metered event. The reservation settles into a Track event named by this + // subtype, so a credit entitlement with neither a resolvable subtype nor a + // positive rate can't be billed and is treated as ungateable. + const resolvedSubtype = eventSubtype ?? entitlement.eventSubtype; + if (!creditId || consumptionRate <= 0 || !resolvedSubtype) { + logger.debug( + `Lease check: flag ${key} credit entitlement is incomplete ` + + `(creditId=${creditId ?? ""}, consumptionRate=${consumptionRate}, subtype=${ + resolvedSubtype ?? "" + }), falling back`, + ); return fallback(); } const creditCost = quantity * consumptionRate; @@ -284,10 +316,16 @@ export async function checkWithLease( const preReservation = postReserveBalance + creditCost; const substituted = substituteCreditBalance(company, creditId, preReservation); - const wasmOptions = buildPreflightOptions(options); + // Gate precisely on the lease tranche: tell the engine this action costs + // `creditCost` against this credit id. `credit_cost` carries the + // pre-computed `quantity × consumptionRate` — highest precedence for credit + // gates — so the engine checks `preReservation − creditCost ≥ 0`, the same + // arithmetic `tryReserve` just enforced, at the same per-credit-id + // granularity as the lease. + const gatingOptions: CheckFlagOptions = { creditCost: { [creditId]: creditCost } }; let result; try { - result = await datastream.getRulesEngine().checkFlagWithOptions(flag, substituted, user ?? null, wasmOptions); + result = await datastream.getRulesEngine().checkFlagWithOptions(flag, substituted, user ?? null, gatingOptions); } catch (err) { logger.error(`Lease check: WASM eval failed: ${err}`); // Cancel the hold: removes the reservation record and refunds the full @@ -326,34 +364,11 @@ export async function checkWithLease( ); } - // The engine allowed — but keep the hold only if the allow actually came - // from the rule whose credit condition we reserved against. An allow via a - // different rule (company/global override, another plan's rule) grants the - // feature without metering this credit, so keeping the reservation would - // bill usage the matched rule grants for free. Cancel it and return the - // allow without a handle — `trackWithReservation` is never owed. Only a - // definitive mismatch cancels: if the engine doesn't report a ruleId, the - // reservation is kept rather than silently disabling credit gating. - if (result.ruleId !== undefined && result.ruleId !== match.ruleId) { - logger.debug( - `Lease check: flag ${key} allowed by rule ${result.ruleId} (${result.ruleType ?? "unknown type"}), ` + - `not the credit-metered rule ${match.ruleId} — cancelling reservation, no credits billed`, - ); - await cancelReservation(reservations, reservation, logger); - return emitFlagCheck( - deps, - evalCtx, - { - allowed: true, - value: true, - reason: result.reason ?? "allowed_without_credit_rule", - entitlement: normalizeEntitlement(result.entitlement), - flagKey: result.flagKey ?? key, - flagId: result.flagId, - }, - ids, - ); - } + // The engine allowed against the substituted lease balance — keep the hold. + // No rule-match disambiguation is needed: the entitlement probe already + // established this company's matched entitlement is the credit one, so an + // override/boolean grant would have skipped the reserve path above. Bumping + // only the credit balance can't make a different rule match here. // Fire-and-forget low-water-mark refresh now that we've debited. void manager.maybeExtendInBackground(company.id, creditId); @@ -402,109 +417,6 @@ export function buildPreflightOptions(options: CheckOptions): CheckFlagOptions | return { usage: options.usage }; } -interface CreditConditionMatch { - condition: api.RulesengineCondition & { creditId: string }; - /** - * Id of the rule the condition belongs to. After the gating eval, the - * reservation is kept only if the engine's matched `ruleId` is this rule — - * an allow via a different rule (company/global override, another plan's - * rule) grants the feature without metering this credit, so the hold is - * cancelled instead of billed. - */ - ruleId: 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, ruleId: rule.id }; - for (const group of rule.conditionGroups ?? []) { - const inGroup = matchInConditions(group.conditions ?? [], eventSubtype, creditId); - if (inGroup) return { ...inGroup, ruleId: rule.id }; - } - } - return null; -} - -function matchInConditions( - conditions: api.RulesengineCondition[], - eventSubtype: string | undefined, - creditId?: string, -): Omit | 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, diff --git a/src/credits/lease-manager.ts b/src/credits/lease-manager.ts index 41472dbd..03b5a1dd 100644 --- a/src/credits/lease-manager.ts +++ b/src/credits/lease-manager.ts @@ -88,9 +88,10 @@ export class CreditLeaseManager { // 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. + // every path that acts on a lease (`tryReserve`, and the expiry check + // just above) re-guards on expiry before trusting it. 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); diff --git a/src/credits/types.ts b/src/credits/types.ts index 3e5158bc..bc318c66 100644 --- a/src/credits/types.ts +++ b/src/credits/types.ts @@ -167,15 +167,17 @@ export interface CheckResult { /** * 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. + * For a credit-metered feature the entitlement carries the lease-aware + * three-way split the server computes: `creditRemaining` (drawable now, open + * lease holds excluded), `creditReserved` (the open lease's unspent hold), + * and `creditSettled` (`creditRemaining + creditReserved` = the balance net + * of actual consumption). Bind a user-facing "credits remaining" counter to + * `creditSettled` — it's invariant to in-flight lease holds, so it doesn't + * dip mid-call the way the raw server balance does. + * + * These come from the company's resolved entitlement, so the lease-balance + * substitution a lease-bearing `check()` applies for gating does NOT distort + * them — the same values surface whether or not `usage` was passed. */ entitlement?: api.RulesengineFeatureEntitlement; /** Flag key checked. */ diff --git a/src/rules-engine.ts b/src/rules-engine.ts index 737276f9..f0eba79a 100644 --- a/src/rules-engine.ts +++ b/src/rules-engine.ts @@ -10,13 +10,17 @@ export interface WasmFeatureEntitlement { softLimit?: number; usage?: number; eventName?: string; + eventSubtype?: string; metricPeriod?: Schematic.RulesengineMetricPeriod; monthReset?: Schematic.RulesengineMetricPeriodMonthReset; metricResetAt?: string; creditId?: string; + consumptionRate?: number; creditTotal?: number; creditUsed?: number; creditRemaining?: number; + creditReserved?: number; + creditSettled?: number; } /** Result returned by the WASM rules engine */ diff --git a/src/wrapper.ts b/src/wrapper.ts index 2359f84a..6ba90dda 100644 --- a/src/wrapper.ts +++ b/src/wrapper.ts @@ -987,81 +987,6 @@ export class SchematicClient extends BaseClient { }); } - /** - * Lease-aware credit balance for display, in the same three-way split the - * server exposes on `listCompanyCreditBalances`: - * - * - `remaining`: credits available to draw right now — the server balance - * (`company.creditBalances[creditId]`, already net of any active lease - * tranche): `B − L`. - * - `reserved`: the active lease's unsettled hold — its unspent local - * balance plus open reservations (`L − spent`), reported as the - * aggregate the server models. - * - `settled`: the economic balance the end user should see — - * `remaining + reserved = B − spent`. Invariant to in-flight holds, so - * a counter bound to it 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 balance. - * - * `reserved` is 0 when no lease is held or the held lease has expired — the - * server refunds an expired lease's unspent hold into the balance, so - * counting a stale local entry 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 `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; - // Display path: tolerate an unreachable lease store by reporting no - // local hold rather than throwing. - const leaseEntry = await Promise.resolve() - .then(() => this.leaseStore?.get(snapshot.id, creditId)) - .catch((err) => { - this.logger.debug(`getCreditBalance: lease store read failed (${err})`); - return undefined; - }); - // 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. diff --git a/tests/unit/credits/check-and-track.test.ts b/tests/unit/credits/check-and-track.test.ts index 6497ab90..1172f8b3 100644 --- a/tests/unit/credits/check-and-track.test.ts +++ b/tests/unit/credits/check-and-track.test.ts @@ -98,35 +98,6 @@ const flag = { ], }; -// 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", @@ -169,6 +140,69 @@ function configureDataStream() { mockDataStream.getCompany.mockResolvedValue(company); } +// The matched credit entitlement the engine probe surfaces. Entitlement-first +// resolution reads creditId + consumptionRate + eventSubtype straight off this +// (no structural flag scan), so the consumption rate now lives here, not on the +// flag condition. +const CREDIT_ENTITLEMENT = { + featureId: "feat", + featureKey: "inference", + valueType: "credit", + creditId: "bilcr_inference", + consumptionRate: 10, + eventSubtype: "inference_tokens", + creditTotal: 1000, + creditUsed: 0, + creditRemaining: 1000, +}; + +// The lease path calls the rules engine twice: first an unsubstituted +// entitlement probe (`options === null`) to identify the matched entitlement, +// then a substituted gating eval (`options.creditCost` set) once a reservation +// is held. This wires both: `probe`/`gate` override the respective result, and +// the probe defaults to the credit entitlement above so the path proceeds to +// reserve. The fail-open re-eval in `handleLeaseFailure` also runs without +// `creditCost`, so it is served by the probe branch too. +function configureEngine({ + probe = {}, + gate = {}, +}: { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + probe?: Record; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + gate?: Record; +} = {}) { + mockRulesEngine.checkFlagWithOptions.mockImplementation( + // eslint-disable-next-line @typescript-eslint/no-explicit-any + (_flag: any, _co: any, _user: any, options: any) => { + if (options && options.creditCost) { + return Promise.resolve({ + value: true, + reason: "matched", + flagKey: "inference", + flagId: "flag_1", + ruleId: "rule_1", + entitlement: CREDIT_ENTITLEMENT, + ...gate, + }); + } + return Promise.resolve({ + value: true, + reason: "probe", + flagKey: "inference", + flagId: "flag_1", + entitlement: CREDIT_ENTITLEMENT, + ...probe, + }); + }, + ); +} + +// The substituted gating eval, found by its `creditCost` options envelope. +function gateCall() { + return mockRulesEngine.checkFlagWithOptions.mock.calls.find((c) => c[3] && c[3].creditCost); +} + function makeClient() { return new SchematicClient({ apiKey: "test-key", @@ -199,21 +233,7 @@ 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, - }, - }); + configureEngine(); const client = makeClient(); const result = await client.check({ company: { id: "co_1" } }, "inference", { @@ -227,30 +247,19 @@ describe("client.check (lease path)", () => { 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: { eventSubtype: "inference_tokens", quantity: 50 } }); + // The gating eval sees the substituted pre-reservation balance (1000) and + // the pre-computed credit cost (50 × 10). + const gate = gateCall(); + expect(gate?.[1].creditBalances.bilcr_inference).toBe(1000); + expect(gate?.[3]).toEqual({ creditCost: { bilcr_inference: 500 } }); 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); - const mixedCompany = { - ...company, - creditBalances: { bilcr_usd: 0, bilcr_ai: 5000 }, - }; - mockDataStream.getCachedCompany.mockResolvedValue(mixedCompany); - mockDataStream.getCompany.mockResolvedValue(mixedCompany); + it("leases the credit the company's matched plan entitlement uses", async () => { + // The entitlement probe names the company's matched plan credit + // (AI credits, rate 5) regardless of flag shape — we lease and meter + // that, reading both creditId and consumptionRate straight off it. + configureDataStream(); mockAcquireCreditLease.mockResolvedValue({ data: { id: "lse_ai", @@ -263,20 +272,13 @@ describe("client.check (lease path)", () => { }, 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 aiEntitlement = { + ...CREDIT_ENTITLEMENT, + creditId: "bilcr_ai", + consumptionRate: 5, + creditRemaining: 5000, + }; + configureEngine({ probe: { entitlement: aiEntitlement }, gate: { entitlement: aiEntitlement } }); const client = makeClient(); const result = await client.check({ company: { id: "co_1" } }, "inference", { @@ -285,7 +287,6 @@ describe("client.check (lease path)", () => { }); 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); @@ -293,23 +294,24 @@ describe("client.check (lease path)", () => { 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); + expect(gateCall()?.[3]).toEqual({ creditCost: { bilcr_ai: 250 } }); await client.close(); }); - it("cancels the reservation when the allow comes from a different (non-credit) rule", async () => { - // The company is granted the feature by an override rule, not the - // credit-metered rule we reserved against. Billing the reservation - // would charge credits for usage the override grants for free — the - // hold must be cancelled and the allow returned without a handle. + it("skips the lease entirely when the matched entitlement is not credit-metered", async () => { + // A company/global override (or any boolean/numeric entitlement) grants + // the feature without drawing a credit. The probe surfaces a non-credit + // entitlement, so the lease path never acquires or reserves — it defers + // to the plain check, which returns the allow without a handle. configureSuccessfulAcquire(); configureDataStream(); - mockRulesEngine.checkFlagWithOptions.mockResolvedValueOnce({ + configureEngine({ + probe: { value: true, entitlement: { featureId: "feat", featureKey: "inference", valueType: "boolean" } }, + }); + mockDataStream.checkFlag.mockResolvedValue({ value: true, reason: "company override", flagKey: "inference", - flagId: "flag_1", - ruleId: "rule_override", - ruleType: "company_override", }); const client = makeClient(); @@ -321,37 +323,16 @@ describe("client.check (lease path)", () => { expect(result.allowed).toBe(true); expect(result.value).toBe(true); expect(result.reservation).toBeUndefined(); - - // The cancelled hold was refunded: a follow-up check that DOES match - // the credit rule sees the full lease balance (1000, not 1000-500) - // substituted into its evaluation. - mockRulesEngine.checkFlagWithOptions.mockResolvedValueOnce({ - value: true, - reason: "matched", - flagKey: "inference", - flagId: "flag_1", - ruleId: "rule_1", - }); - const second = await client.check({ company: { id: "co_1" } }, "inference", { - usage: 50, - eventSubtype: "inference_tokens", - }); - expect(second.reservation?.creditsReserved).toBe(500); - const secondCallArgs = mockRulesEngine.checkFlagWithOptions.mock.calls[1]; - expect(secondCallArgs[1].creditBalances.bilcr_inference).toBe(1000); + // No reserve-then-cancel round-trip: the lease is never acquired. + expect(mockAcquireCreditLease).not.toHaveBeenCalled(); + expect(gateCall()).toBeUndefined(); await client.close(); }); - it("keeps the reservation when the engine reports the credit rule's ruleId", async () => { + it("keeps the reservation when the gating eval allows", async () => { configureSuccessfulAcquire(); configureDataStream(); - mockRulesEngine.checkFlagWithOptions.mockResolvedValue({ - value: true, - reason: "matched", - flagKey: "inference", - flagId: "flag_1", - ruleId: "rule_1", - }); + configureEngine({ gate: { value: true } }); const client = makeClient(); const result = await client.check({ company: { id: "co_1" } }, "inference", { @@ -364,14 +345,10 @@ describe("client.check (lease path)", () => { await client.close(); }); - it("returns allowed=false and refunds reservation when engine denies", async () => { + it("returns allowed=false and refunds reservation when the gating eval denies", async () => { configureSuccessfulAcquire(); configureDataStream(); - mockRulesEngine.checkFlagWithOptions.mockResolvedValue({ - value: false, - reason: "denied", - flagKey: "inference", - }); + configureEngine({ gate: { value: false, reason: "denied" } }); const client = makeClient(); const result = await client.check({ company: { id: "co_1" } }, "inference", { @@ -387,6 +364,7 @@ describe("client.check (lease path)", () => { it("fails closed when lease acquire errors and onAcquireFailure='fail-closed'", async () => { configureFailingAcquire(); configureDataStream(); + configureEngine(); const client = makeClient(); const result = await client.check({ company: { id: "co_1" } }, "inference", { @@ -397,13 +375,16 @@ describe("client.check (lease path)", () => { expect(result.allowed).toBe(false); expect(result.reservation).toBeUndefined(); - expect(mockRulesEngine.checkFlagWithOptions).not.toHaveBeenCalled(); + // The entitlement probe ran (to resolve the credit), but no gating eval — + // acquire failed before any reservation. + expect(gateCall()).toBeUndefined(); await client.close(); }); it("defaults to fail-closed when onAcquireFailure is not specified", async () => { configureFailingAcquire(); configureDataStream(); + configureEngine(); const client = makeClient(); const result = await client.check({ company: { id: "co_1" } }, "inference", { @@ -413,19 +394,14 @@ describe("client.check (lease path)", () => { expect(result.allowed).toBe(false); expect(result.reservation).toBeUndefined(); - expect(mockRulesEngine.checkFlagWithOptions).not.toHaveBeenCalled(); + expect(gateCall()).toBeUndefined(); await client.close(); }); it("respects explicit fail-open override on acquire failure, still evaluating the rules", async () => { configureFailingAcquire(); configureDataStream(); - mockRulesEngine.checkFlagWithOptions.mockResolvedValue({ - value: true, - reason: "matched", - flagKey: "inference", - flagId: "flag_1", - }); + configureEngine(); const client = makeClient(); const result = await client.check({ company: { id: "co_1" } }, "inference", { @@ -439,24 +415,21 @@ describe("client.check (lease path)", () => { expect(result.err).toBe("lease_acquire_failed"); // Fail-open still runs the engine — with the credit balance substituted // to an effectively unlimited value so only the credit gate is bypassed. - expect(mockRulesEngine.checkFlagWithOptions).toHaveBeenCalledTimes(1); - const callArgs = mockRulesEngine.checkFlagWithOptions.mock.calls[0]; - expect(callArgs[1].creditBalances.bilcr_inference).toBe(Number.MAX_SAFE_INTEGER); - expect(callArgs[3]).toEqual({ eventUsage: { eventSubtype: "inference_tokens", quantity: 50 } }); + const failOpenCall = mockRulesEngine.checkFlagWithOptions.mock.calls.find( + (c) => c[1].creditBalances.bilcr_inference === Number.MAX_SAFE_INTEGER, + ); + expect(failOpenCall).toBeDefined(); + expect(failOpenCall?.[3]).toEqual({ eventUsage: { eventSubtype: "inference_tokens", quantity: 50 } }); await client.close(); }); it("fail-open still denies a company the rules do not entitle", async () => { configureFailingAcquire(); configureDataStream(); - // The engine denies for a non-credit reason (e.g. no matching plan rule) - // even with the substituted unlimited balance. - mockRulesEngine.checkFlagWithOptions.mockResolvedValue({ - value: false, - reason: "no matching rule", - flagKey: "inference", - flagId: "flag_1", - }); + // The fail-open re-eval denies for a non-credit reason (e.g. no matching + // plan rule) even with the substituted unlimited balance. The probe + // still resolves the credit so the path reaches acquire. + configureEngine({ probe: { value: false, reason: "no matching rule" } }); const client = makeClient(); const result = await client.check({ company: { id: "co_1" } }, "inference", { @@ -471,10 +444,25 @@ describe("client.check (lease path)", () => { await client.close(); }); - it("fail-open falls back to a blanket allow when the fail-open evaluation itself errors", async () => { + it("fail-open falls back to a blanket allow when the fail-open eval itself errors", async () => { + // Probe resolves the credit (so the path reaches acquire), acquire fails, + // then the fail-open re-eval in handleLeaseFailure throws — leaving only + // the static fail-open blanket allow. configureFailingAcquire(); configureDataStream(); - mockRulesEngine.checkFlagWithOptions.mockRejectedValue(new Error("wasm exploded")); + mockRulesEngine.checkFlagWithOptions.mockImplementation( + // eslint-disable-next-line @typescript-eslint/no-explicit-any + (_flag: any, _co: any, _user: any, options: any) => { + if (options) return Promise.reject(new Error("wasm exploded")); // fail-open re-eval + return Promise.resolve({ + value: true, + reason: "probe", + flagKey: "inference", + flagId: "flag_1", + entitlement: CREDIT_ENTITLEMENT, + }); + }, + ); const client = makeClient(); const result = await client.check({ company: { id: "co_1" } }, "inference", { @@ -488,6 +476,27 @@ describe("client.check (lease path)", () => { await client.close(); }); + it("falls back to the plain check when the entitlement probe itself errors", async () => { + // A probe failure is a resolution miss, not the gate — defer to the plain + // check (with its own degradation) rather than hard-denying. No lease is + // acquired. + configureSuccessfulAcquire(); + configureDataStream(); + mockRulesEngine.checkFlagWithOptions.mockRejectedValue(new Error("wasm exploded")); + mockDataStream.checkFlag.mockResolvedValue({ value: true, reason: "fallback", flagKey: "inference" }); + + const client = makeClient(); + const result = await client.check({ company: { id: "co_1" } }, "inference", { + usage: 50, + eventSubtype: "inference_tokens", + }); + + expect(result.allowed).toBe(true); + expect(result.reservation).toBeUndefined(); + expect(mockAcquireCreditLease).not.toHaveBeenCalled(); + await client.close(); + }); + it("falls back to plain check when no preflight options are provided", async () => { configureDataStream(); mockCheckFlag.mockResolvedValue({ @@ -578,22 +587,16 @@ describe("client.check (lease path)", () => { await client.close(); }); - it("falls back when neither the caller nor the condition supplies an event subtype", async () => { + it("falls back when neither the caller nor the entitlement supplies an event subtype", async () => { // The reservation settles into a Track event named by the subtype; an // empty event name would consume the reservation while billing - // nothing, so the lease path must decline to reserve. + // nothing, so the lease path must decline to reserve when the matched + // credit entitlement carries no subtype and the caller named none. configureSuccessfulAcquire(); configureDataStream(); - const noSubtypeFlag = { - ...flag, - rules: [ - { - ...flag.rules[0], - conditions: [{ ...flag.rules[0].conditions[0], eventSubtype: undefined }], - }, - ], - }; - mockDataStream.getFlag.mockResolvedValue(noSubtypeFlag); + configureEngine({ + probe: { entitlement: { ...CREDIT_ENTITLEMENT, eventSubtype: undefined } }, + }); mockDataStream.checkFlag.mockResolvedValue({ value: true, reason: "match", flagKey: "inference" }); const client = makeClient(); @@ -616,12 +619,7 @@ describe("client.check (lease path)", () => { // approving every subsequent reserve (`NaN < x` is always false). configureSuccessfulAcquire(); configureDataStream(); - mockRulesEngine.checkFlagWithOptions.mockResolvedValue({ - value: true, - reason: "matched", - flagKey: "inference", - flagId: "flag_1", - }); + configureEngine(); const client = makeClient(); const denied = await client.check({ company: { id: "co_1" } }, "inference", { @@ -675,14 +673,7 @@ describe("flag_check event reporting on the lease path", () => { it("enqueues a flag_check event when the engine allows", async () => { configureSuccessfulAcquire(); configureDataStream(); - mockRulesEngine.checkFlagWithOptions.mockResolvedValue({ - value: true, - reason: "matched", - flagKey: "inference", - flagId: "flag_1", - ruleId: "rule_1", - companyId: "co_1", - }); + configureEngine({ gate: { value: true, reason: "matched", ruleId: "rule_1", companyId: "co_1" } }); const client = makeClient(); await client.check({ company: { id: "co_1" } }, "inference", { @@ -707,11 +698,7 @@ describe("flag_check event reporting on the lease path", () => { it("enqueues a flag_check event when the engine denies", async () => { configureSuccessfulAcquire(); configureDataStream(); - mockRulesEngine.checkFlagWithOptions.mockResolvedValue({ - value: false, - reason: "denied", - flagKey: "inference", - }); + configureEngine({ gate: { value: false, reason: "denied", companyId: "co_1" } }); const client = makeClient(); await client.check({ company: { id: "co_1" } }, "inference", { @@ -733,6 +720,7 @@ describe("flag_check event reporting on the lease path", () => { it("enqueues a flag_check event on a fail-closed acquire failure", async () => { configureFailingAcquire(); configureDataStream(); + configureEngine(); const client = makeClient(); await client.check({ company: { id: "co_1" } }, "inference", { @@ -827,12 +815,7 @@ 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", - }); + configureEngine(); const client = makeClient(); const res = await client.check({ company: { id: "co_1" } }, "inference", { @@ -902,12 +885,7 @@ describe("client.trackWithReservation", () => { 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", - }); + configureEngine(); const client = makeClient(); const res = await client.check({ company: { id: "co_1" } }, "inference", { usage: 50, @@ -924,12 +902,7 @@ describe("client.trackWithReservation", () => { it("rejects a non-finite or negative actualQuantity without touching the store or billing", async () => { configureSuccessfulAcquire(); configureDataStream(); - mockRulesEngine.checkFlagWithOptions.mockResolvedValue({ - value: true, - reason: "matched", - flagKey: "inference", - flagId: "flag_1", - }); + configureEngine(); const client = makeClient(); const res = await client.check({ company: { id: "co_1" } }, "inference", { usage: 50, @@ -968,12 +941,7 @@ describe("evalCtx user resolution on the lease path", () => { configureSuccessfulAcquire(); configureDataStream(); mockDataStream.getUser.mockResolvedValue(user); - mockRulesEngine.checkFlagWithOptions.mockResolvedValue({ - value: true, - reason: "matched", - flagKey: "inference", - flagId: "flag_1", - }); + configureEngine(); const client = makeClient(); const result = await client.check({ company: { id: "co_1" }, user: { id: "user_1" } }, "inference", { @@ -983,7 +951,7 @@ describe("evalCtx user resolution on the lease path", () => { expect(result.allowed).toBe(true); expect(mockDataStream.getUser).toHaveBeenCalledWith({ id: "user_1" }); - // The resolved user is threaded into the gating eval — user-targeted + // The resolved user is threaded into both engine calls — user-targeted // rules and overrides apply on the lease path. const callArgs = mockRulesEngine.checkFlagWithOptions.mock.calls[0]; expect(callArgs[2]).toBe(user); @@ -1080,6 +1048,7 @@ describe("store-failure containment (unreachable Redis backend)", () => { it("check() resolves fail-closed instead of rejecting when the store is down", async () => { configureSuccessfulAcquire(); configureDataStream(); + configureEngine(); const client = makeRedisBackedClient(); const result = await client.check({ company: { id: "co_1" } }, "inference", { @@ -1096,12 +1065,7 @@ describe("store-failure containment (unreachable Redis backend)", () => { it("check() honors fail-open (rules still evaluated) when the store is down", async () => { configureSuccessfulAcquire(); configureDataStream(); - mockRulesEngine.checkFlagWithOptions.mockResolvedValue({ - value: true, - reason: "matched", - flagKey: "inference", - flagId: "flag_1", - }); + configureEngine(); const client = makeRedisBackedClient(); const result = await client.check({ company: { id: "co_1" } }, "inference", { @@ -1112,8 +1076,10 @@ describe("store-failure containment (unreachable Redis backend)", () => { expect(result.allowed).toBe(true); expect(result.reservation).toBeUndefined(); - const callArgs = mockRulesEngine.checkFlagWithOptions.mock.calls[0]; - expect(callArgs[1].creditBalances.bilcr_inference).toBe(Number.MAX_SAFE_INTEGER); + const failOpenCall = mockRulesEngine.checkFlagWithOptions.mock.calls.find( + (c) => c[1].creditBalances.bilcr_inference === Number.MAX_SAFE_INTEGER, + ); + expect(failOpenCall).toBeDefined(); await client.close(); }); @@ -1150,12 +1116,7 @@ describe("close() lease release", () => { configureSuccessfulAcquire(); configureDataStream(); mockReleaseCreditLease.mockResolvedValue({}); - mockRulesEngine.checkFlagWithOptions.mockResolvedValue({ - value: true, - reason: "matched", - flagKey: "inference", - flagId: "flag_1", - }); + configureEngine(); const client = makeClient(); await client.check({ company: { id: "co_1" } }, "inference", { usage: 50, eventSubtype: "inference_tokens" }); @@ -1172,12 +1133,7 @@ describe("per-check timeout threading", () => { it("threads timeoutMs to the lease acquire wire call", async () => { configureSuccessfulAcquire(); configureDataStream(); - mockRulesEngine.checkFlagWithOptions.mockResolvedValue({ - value: true, - reason: "matched", - flagKey: "inference", - flagId: "flag_1", - }); + configureEngine(); const client = makeClient(); await client.check({ company: { id: "co_1" } }, "inference", { diff --git a/tests/unit/credits/wasm-credit-gate.test.ts b/tests/unit/credits/wasm-credit-gate.test.ts index 41169f14..88fd055f 100644 --- a/tests/unit/credits/wasm-credit-gate.test.ts +++ b/tests/unit/credits/wasm-credit-gate.test.ts @@ -2,12 +2,13 @@ * 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 + * contract that actually matters most — resolving the matched credit + * entitlement from the probe, substituting the lease balance into the company's + * `creditBalances`, and letting the WASM's `credit_cost` 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 (`credit_cost` / `event_usage`) or in the * camelCase entity shape the SDK now feeds the engine would fail here instead of * silently mis-gating in production. */ @@ -87,13 +88,30 @@ function creditFlag(extraConditions: object[] = []): api.RulesengineFlag { } as unknown as api.RulesengineFlag; } -function company(creditBalance: number): api.RulesengineCompany { +// The company carries its resolved credit entitlement for the feature — the +// shape the datastream cache holds after a plan assignment. Entitlement-first +// resolution reads creditId / consumptionRate / eventSubtype off this via the +// engine probe, so it must be present for the lease path to gate on the credit. +function company(creditBalance: number, entitlements?: object[]): api.RulesengineCompany { return { id: "co", accountId: "acct", environmentId: "env", keys: {}, creditBalances: { [CREDIT_ID]: creditBalance }, + entitlements: entitlements ?? [ + { + featureId: "feat-infer", + featureKey: "infer", + valueType: "credit", + creditId: CREDIT_ID, + consumptionRate: 1, + eventSubtype: EVENT_SUBTYPE, + creditTotal: creditBalance, + creditUsed: 0, + creditRemaining: creditBalance, + }, + ], } as unknown as api.RulesengineCompany; } @@ -210,12 +228,13 @@ describe("credit-lease check against the real WASM engine", () => { expect(reservations.size()).toBe(0); }); - it("cancels the reservation when the WASM allows via an override rule instead of the credit rule", async () => { - // The company is granted the feature by a company_override rule that - // outranks the credit-metered rule. The engine allows — but via a rule - // that doesn't meter this credit, so the reservation taken before the - // eval must be cancelled (no handle, nothing billed) rather than left - // to charge credits for usage the override grants for free. + it("skips the lease entirely when the company's matched entitlement is a (non-credit) override", async () => { + // The company is granted the feature by a company_override that resolves + // to a boolean entitlement — its effective entitlement for the feature is + // no longer credit-metered. The probe surfaces the boolean entitlement, + // so the lease path never acquires or reserves; it defers to the plain + // check. No reserve-then-cancel, and no credits billed for usage the + // override grants for free. const flag = creditFlag(); // eslint-disable-next-line @typescript-eslint/no-explicit-any (flag.rules as any[]).unshift({ @@ -229,21 +248,27 @@ describe("credit-lease check against the real WASM engine", () => { conditions: [companyCondition(["co"])], conditionGroups: [], }); - const { deps, leaseStore, reservations } = makeDeps(engine, flag, company(100), 10_000); + // Effective entitlement after the override: boolean, not credit. + const overrideCompany = company(100, [{ featureId: "feat-infer", featureKey: "infer", valueType: "boolean" }]); + const { deps, leaseStore, reservations } = makeDeps(engine, flag, overrideCompany, 10_000); - const result = await checkWithLease( - deps, - "infer", - evalCtx, - { usage: 50, eventSubtype: EVENT_SUBTYPE }, - failFallback, - ); + let fellBack = false; + const result = await checkWithLease(deps, "infer", evalCtx, { usage: 50, eventSubtype: EVENT_SUBTYPE }, () => { + fellBack = true; + return Promise.resolve({ + allowed: true, + value: true, + reason: "override", + flagKey: "infer", + }); + }); + expect(fellBack).toBe(true); expect(result.allowed).toBe(true); expect(result.value).toBe(true); expect(result.reservation).toBeUndefined(); - // The hold was refunded: lease back to full, reservation table empty. - expect((await leaseStore.get("co", CREDIT_ID))?.localRemainingCredits).toBe(10_000); + // The lease was never acquired or touched. + expect(await leaseStore.get("co", CREDIT_ID)).toBeUndefined(); expect(reservations.size()).toBe(0); }); diff --git a/tests/unit/wrapper.test.ts b/tests/unit/wrapper.test.ts index 640b9512..75c026fc 100644 --- a/tests/unit/wrapper.test.ts +++ b/tests/unit/wrapper.test.ts @@ -293,102 +293,6 @@ describe("SchematicClient wrapper - flag checking behavior", () => { }); }); -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", () => { let consoleSpy: { debug: ReturnType; From a4cbe8f185ddfbe4e3e134f3104cb5134f0b3ddb Mon Sep 17 00:00:00 2001 From: Ben Papillon Date: Mon, 22 Jun 2026 09:55:03 -0700 Subject: [PATCH 2/2] Handle credit_lease_balance partial in datastream merge The API now emits a single credit_lease_balance partial on lease acquire/extend (remaining + reserved together) instead of separate credit_balances and credit_reserved messages. This SDK is lease-aware, gates on remaining locally, and derives no settled balance, so apply remaining exactly like the credit_balances partial and ignore reserved. Without this case the new partial would fall through to the silent unknown-key skip and the cached remaining would go stale, since acquire/extend no longer emit a standalone credit_balances message. --- src/datastream/merge.ts | 23 ++++++++++++++++++++++- tests/unit/datastream/merge.test.ts | 28 ++++++++++++++++++++++++++++ 2 files changed, 50 insertions(+), 1 deletion(-) diff --git a/src/datastream/merge.ts b/src/datastream/merge.ts index 7fd3e670..58349e9d 100644 --- a/src/datastream/merge.ts +++ b/src/datastream/merge.ts @@ -34,7 +34,7 @@ export function deepCopyUser(u: Schematic.RulesengineUser): Schematic.Rulesengin * Partials don't carry refreshed entitlements, so when their derived fields * change in another part of the company we sync them here to match server * behavior: - * - creditRemaining ← credit_balances[credit_id] + * - creditRemaining ← credit_balances[credit_id] (or credit_lease_balance[credit_id].remaining) * - usage ← metric value matching (eventName, metricPeriod, monthReset) * Both are skipped when the partial also sends entitlements wholesale. * @@ -105,6 +105,27 @@ export function partialCompany( updatedBalances = incomingCB; break; } + case "credit_lease_balance": { + // Atomic lease-balance partial: for each credit it carries the + // post-move remaining together with the open lease hold + // (reserved), emitted as a single message on lease + // acquire/extend so a consumer never derives state from a + // remaining and a reserved captured at different instants. This + // SDK is lease-aware, gates on remaining locally, and derives no + // settled balance, so we apply remaining exactly like the + // credit_balances partial and ignore reserved. + const existingCB = (merged.creditBalances ?? {}) as Record; + const incoming = (partial[key] ?? {}) as Record; + const remainingByCredit: Record = {}; + for (const [creditId, lb] of Object.entries(incoming)) { + if (lb && typeof lb.remaining === "number") { + remainingByCredit[creditId] = lb.remaining; + } + } + merged.creditBalances = { ...existingCB, ...remainingByCredit }; + updatedBalances = { ...(updatedBalances ?? {}), ...remainingByCredit }; + break; + } case "metrics": { const existingMetrics = (merged.metrics ?? []) as Schematic.RulesengineCompanyMetric[]; const incomingMetrics = ((partial[key] ?? []) as unknown[]).map((m) => diff --git a/tests/unit/datastream/merge.test.ts b/tests/unit/datastream/merge.test.ts index 017142eb..ad751fb3 100644 --- a/tests/unit/datastream/merge.test.ts +++ b/tests/unit/datastream/merge.test.ts @@ -196,6 +196,34 @@ describe('partialCompany', () => { expect(origEnts[0].creditRemaining).toBe(100.0); }); + test('credit_lease_balance applies remaining (ignoring reserved) and re-derives entitlement creditRemaining', () => { + const existing = baseCompany(); + (existing as unknown as Record).entitlements = [ + { featureId: 'feat-1', featureKey: 'feature-one', valueType: 'credit', creditId: 'credit-1', creditRemaining: 100.0 }, + { featureId: 'feat-2', featureKey: 'feature-two', valueType: 'credit', creditId: 'credit-2', creditRemaining: 0 }, + ]; + + // Atomic lease-balance partial: remaining + reserved for credit-1, no entitlements. + const partial = { id: 'co-1', credit_lease_balance: { 'credit-1': { remaining: 42.0, reserved: 58.0 } } }; + + const merged = partialCompany(existing, partial); + const m = merged as unknown as Record; + + // Company-level balance updated to remaining; reserved is ignored. + expect(m.creditBalances).toEqual({ 'credit-1': 42.0 }); + + const ents = m.entitlements as Record[]; + // credit-1 entitlement re-derived from incoming remaining + expect(ents[0].creditRemaining).toBe(42.0); + expect(ents[0].credit_remaining).toBeUndefined(); + // credit-2 not in the partial, left untouched + expect(ents[1].creditRemaining).toBe(0); + + // Original not mutated + const origEnts = (existing as unknown as Record).entitlements as Record[]; + expect(origEnts[0].creditRemaining).toBe(100.0); + }); + test('metrics update re-derives entitlement usage', () => { const existing = baseCompany(); (existing as unknown as Record).metrics = [