From 9e13faf45b7edb43a36a3454022f8d022bd9a371 Mon Sep 17 00:00:00 2001 From: Ben Papillon Date: Thu, 8 Oct 2026 17:39:59 -0700 Subject: [PATCH] pass every preflight field to the rules engine --- .../RulesEngine/Wasm/EventQuantitiesTests.cs | 280 ++++++++++++++++++ .../Wasm/testdata/quantity_cost.json | 185 ++++++++++++ .../SchematicHQ.Client.Test.Custom.props | 3 + .../RulesEngine/Wasm/WasmRulesEngine.cs | 37 ++- 4 files changed, 494 insertions(+), 11 deletions(-) create mode 100644 src/SchematicHQ.Client.Test/RulesEngine/Wasm/EventQuantitiesTests.cs create mode 100644 src/SchematicHQ.Client.Test/RulesEngine/Wasm/testdata/quantity_cost.json diff --git a/src/SchematicHQ.Client.Test/RulesEngine/Wasm/EventQuantitiesTests.cs b/src/SchematicHQ.Client.Test/RulesEngine/Wasm/EventQuantitiesTests.cs new file mode 100644 index 00000000..f76dd4de --- /dev/null +++ b/src/SchematicHQ.Client.Test/RulesEngine/Wasm/EventQuantitiesTests.cs @@ -0,0 +1,280 @@ +using System.Text.Json; +using System.Text.Json.Nodes; +using Microsoft.Extensions.Logging.Abstractions; +using NUnit.Framework; +using SchematicHQ.Client; +using SchematicHQ.Client.Core; +using SchematicHQ.Client.RulesEngine; + +namespace SchematicHQ.Client.Test.RulesEngine.Wasm +{ + /// + /// The event_quantities preflight and the quantity_rates it prices from, run through the + /// WASM engine. Flags, companies and preflights are read from snake_case JSON, as DataStream + /// and the API send them, so a field the generated models would drop fails here. + /// + /// The generated has no event_quantities property + /// until the API's preflight body carries it and Fern regenerates; until then it rides in the + /// body's extension data, which is enough to pin that the engine path forwards it. + /// + [TestFixture] + public class EventQuantitiesTests + { + private const string Subtype = "chat"; + private const string CreditId = "credit-abc"; + private static readonly JsonObject Rates = new() { ["input_tokens"] = 0.001, ["output_tokens"] = 0.01 }; + + // 1 request x 0.5 + (1000 - 400 cached) x 0.001 + 100 x 0.01 = 2.1. The cached tokens are + // unrated, so they cost nothing but still come out of input. + private static readonly JsonObject Quantities = new() + { + ["input_tokens"] = 1000, + ["cached_input_tokens"] = 400, + ["output_tokens"] = 100, + }; + + private WasmRulesEngine _engine = null!; + + [OneTimeSetUp] + public void SetUp() + { + _engine = new WasmRulesEngine(NullLogger.Instance); + _engine.Initialize(); + } + + [OneTimeTearDown] + public void TearDown() => _engine.Dispose(); + + /// + /// A single credit-balance rule priced like an inference entitlement: requests at the + /// consumption rate, tokens at their own rates. + /// + private static RulesengineFlag InferenceFlag(double consumptionRate, JsonNode? quantityRates) + { + var condition = new JsonObject + { + ["id"] = "cond-1", + ["account_id"] = "acct", + ["environment_id"] = "env", + ["condition_type"] = "credit", + ["operator"] = "lt", + ["resource_ids"] = new JsonArray(), + ["trait_value"] = "", + ["credit_id"] = CreditId, + ["consumption_rate"] = consumptionRate, + ["event_subtype"] = Subtype, + }; + if (quantityRates != null) + { + condition["quantity_rates"] = quantityRates.DeepClone(); + } + var flag = new JsonObject + { + ["id"] = "flag-1", + ["account_id"] = "acct", + ["environment_id"] = "env", + ["key"] = "chat", + ["default_value"] = false, + ["rules"] = new JsonArray + { + new JsonObject + { + ["id"] = "rule-1", + ["account_id"] = "acct", + ["environment_id"] = "env", + ["name"] = "Credits", + ["rule_type"] = "plan_entitlement", + ["priority"] = 0, + ["value"] = true, + ["conditions"] = new JsonArray { condition }, + ["condition_groups"] = new JsonArray(), + }, + }, + }; + return JsonUtils.Deserialize(flag.ToJsonString()); + } + + private static JsonObject CompanyNode(double balance) => + new() + { + ["id"] = "co", + ["account_id"] = "acct", + ["environment_id"] = "env", + ["keys"] = new JsonObject { ["id"] = "co" }, + ["credit_balances"] = new JsonObject { [CreditId] = balance }, + ["billing_product_ids"] = new JsonArray(), + ["crm_product_ids"] = new JsonArray(), + ["plan_ids"] = new JsonArray(), + ["plan_version_ids"] = new JsonArray(), + ["metrics"] = new JsonArray(), + ["traits"] = new JsonArray(), + ["rules"] = new JsonArray(), + }; + + private static RulesengineCompany Company(double balance) => + JsonUtils.Deserialize(CompanyNode(balance).ToJsonString()); + + private static PreflightRequestBody Preflight(JsonObject options) => + JsonUtils.Deserialize(options.ToJsonString()); + + private static PreflightRequestBody EventQuantities(string subtype, double? quantity, JsonNode? quantities) + { + var eventQuantities = new JsonObject { ["event_subtype"] = subtype }; + if (quantity.HasValue) + { + eventQuantities["quantity"] = quantity.Value; + } + if (quantities != null) + { + eventQuantities["quantities"] = quantities.DeepClone(); + } + return Preflight(new JsonObject { ["event_quantities"] = eventQuantities }); + } + + private CheckFlagResult Check(double balance, PreflightRequestBody preflight) => + _engine.CheckFlag(Company(balance), null, InferenceFlag(0.5, Rates), preflight); + + [Test] + public void PassesWhenTheBalanceCoversTheCall() + { + var result = Check(2.1, EventQuantities(Subtype, null, Quantities)); + + Assert.That(result.RuleId, Is.EqualTo("rule-1")); + Assert.That(result.Value, Is.True); + } + + [Test] + public void RefusesWhenTheBalanceFallsShort() + { + // Covers the request and the legacy single unit, not the tokens. + var result = Check(2.0, EventQuantities(Subtype, null, Quantities)); + + Assert.That(result.RuleId, Is.Null); + Assert.That(result.Value, Is.False); + } + + [Test] + public void IgnoredForAnotherSubtype() + { + var other = EventQuantities("other", null, new JsonObject { ["input_tokens"] = 1e6 }); + + Assert.That(Check(1.0, other).RuleId, Is.EqualTo("rule-1")); + } + + [Test] + public void QuantityMultipliesTheBaseNotTheQuantities() + { + // 3 x 0.5 + 600 x 0.001 + 100 x 0.01 = 3.1. + var preflight = EventQuantities(Subtype, 3, Quantities); + + Assert.That(Check(3.1, preflight).RuleId, Is.EqualTo("rule-1")); + Assert.That(Check(3.0, preflight).RuleId, Is.Null); + } + + [Test] + public void EventQuantitiesBeatsEventUsage() + { + // event_usage alone would ask 1 x 0.5, which 2.0 covers. + var preflight = Preflight( + new JsonObject + { + ["event_usage"] = new JsonObject { ["event_subtype"] = Subtype, ["quantity"] = 1 }, + ["event_quantities"] = new JsonObject + { + ["event_subtype"] = Subtype, + ["quantities"] = Quantities.DeepClone(), + }, + }); + + Assert.That(Check(2.0, preflight).RuleId, Is.Null); + } + + [Test] + public void AnEmptyCreditCostStillDeclaresNothing() + { + var result = _engine.CheckFlag( + Company(1.0), + null, + InferenceFlag(0.5, Rates), + new PreflightRequestBody { CreditCost = new Dictionary() }); + + // 1.0 covers the legacy single-unit check at 0.5. + Assert.That(result.RuleId, Is.EqualTo("rule-1")); + } + + [Test] + public void ACompanyEntitlementsQuantityRatesReachTheResult() + { + var company = CompanyNode(0); + company["entitlements"] = new JsonArray + { + new JsonObject + { + ["feature_id"] = "feat-1", + ["feature_key"] = "chat", + ["value_type"] = "credit", + ["quantity_rates"] = Rates.DeepClone(), + }, + }; + + var result = _engine.CheckFlag( + JsonUtils.Deserialize(company.ToJsonString()), + null, + InferenceFlag(0.5, null)); + + Assert.That(result.Entitlement, Is.Not.Null); + var rates = result.Entitlement!.AdditionalProperties["quantity_rates"]; + Assert.That(rates.GetProperty("input_tokens").GetDouble(), Is.EqualTo(0.001)); + Assert.That(rates.GetProperty("output_tokens").GetDouble(), Is.EqualTo(0.01)); + } + + private static IEnumerable FixtureCases() + { + // quantity_cost.json is copied verbatim from schematic-api's + // api/lib/rulesengine/testdata/quantity_cost.json. The API's burn and the engine both + // price every case there; running them through the WASM engine here pins that this + // SDK's wire shape for quantity_rates and event_quantities reaches that pricing intact. + var path = Path.Combine( + TestContext.CurrentContext.TestDirectory, + "RulesEngine", + "Wasm", + "testdata", + "quantity_cost.json"); + var fixture = JsonNode.Parse(File.ReadAllText(path))!; + foreach (var tc in fixture["cases"]!.AsArray()) + { + var quantity = tc!["quantity"]?.GetValue(); + var quantities = tc["quantities"]?.AsObject(); + // The preflight rejects negative quantities before pricing. + if (quantity < 0 || (quantities?.Any(q => q.Value!.GetValue() < 0) ?? false)) + { + continue; + } + yield return new TestCaseData(tc.ToJsonString()).SetName(tc["name"]!.GetValue()); + } + } + + [TestCaseSource(nameof(FixtureCases))] + public void SharedQuantityCostFixture(string caseJson) + { + var tc = JsonNode.Parse(caseJson)!; + var flag = InferenceFlag(tc["consumption_rate"]!.GetValue(), tc["quantity_rates"]); + var preflight = EventQuantities(Subtype, tc["quantity"]?.GetValue(), tc["quantities"]); + var cost = tc["expected_cost"]!.GetValue(); + + // A cost priced to zero gates on balance > 0, so the smallest positive balance passes + // and zero does not. + var covers = cost * (1 + 1e-9) + 1e-9; + var shortOf = cost == 0 ? 0 : cost * (1 - 1e-6); + + Assert.That( + _engine.CheckFlag(Company(covers), null, flag, preflight).RuleId, + Is.EqualTo("rule-1"), + $"balance {covers} should cover cost {cost}"); + Assert.That( + _engine.CheckFlag(Company(shortOf), null, flag, preflight).RuleId, + Is.Null, + $"balance {shortOf} should not cover cost {cost}"); + } + } +} diff --git a/src/SchematicHQ.Client.Test/RulesEngine/Wasm/testdata/quantity_cost.json b/src/SchematicHQ.Client.Test/RulesEngine/Wasm/testdata/quantity_cost.json new file mode 100644 index 00000000..ba1dba85 --- /dev/null +++ b/src/SchematicHQ.Client.Test/RulesEngine/Wasm/testdata/quantity_cost.json @@ -0,0 +1,185 @@ +{ + "_comment": "Events and the credits they cost: quantity at consumption_rate plus each rated key of quantities at its rate. The API's burn (data.NewCreditBurnBreakdown, tested in apps/entitlements/services) and the rules engine's quantity_cost (rulesengine-rust/src/engine/quantity_cost.rs) both price every case, so the two cannot drift apart. quantity defaults to 1 when absent or zero. quantities are as the event reports them: input_tokens includes the cached and cache-creation subsets.", + "cases": [ + { + "name": "claude_code_golden", + "consumption_rate": 0.5, + "quantity_rates": { + "cache_creation_input_tokens": 0.00125, + "cached_input_tokens": 0.0001, + "input_tokens": 0.001, + "output_tokens": 0.005, + "reasoning_tokens": 0.009 + }, + "quantities": { + "cache_creation_input_tokens": 10979, + "cached_input_tokens": 13628, + "input_tokens": 24617, + "output_tokens": 41 + }, + "expected_cost": 15.80155 + }, + { + "name": "base_quantity_only", + "consumption_rate": 1.5, + "quantity": 3, + "expected_cost": 4.5 + }, + { + "name": "zero_quantity_counts_as_one", + "consumption_rate": 2, + "quantity": 0, + "expected_cost": 2 + }, + { + "name": "base_quantity_does_not_multiply_tokens", + "consumption_rate": 0.5, + "quantity_rates": { + "input_tokens": 1 + }, + "quantity": 3, + "quantities": { + "input_tokens": 2000 + }, + "expected_cost": 2001.5 + }, + { + "name": "cache_subsets_overrunning_input_clamp_at_zero", + "consumption_rate": 0, + "quantity_rates": { + "cache_creation_input_tokens": 1, + "cached_input_tokens": 1, + "input_tokens": 1, + "output_tokens": 1 + }, + "quantities": { + "cache_creation_input_tokens": 300, + "cached_input_tokens": 400, + "input_tokens": 500 + }, + "expected_cost": 700 + }, + { + "name": "quantities_clamp_at_the_cap", + "consumption_rate": 0, + "quantity_rates": { + "input_tokens": 1, + "output_tokens": 1, + "reasoning_tokens": 1 + }, + "quantities": { + "input_tokens": 2000000000, + "output_tokens": 2000000000, + "reasoning_tokens": 2000000000 + }, + "expected_cost": 3000000000 + }, + { + "name": "zero_rates_and_zero_counts_price_nothing", + "consumption_rate": 0, + "quantity_rates": { + "cached_input_tokens": 3, + "input_tokens": 0.25, + "output_tokens": 0 + }, + "quantities": { + "cached_input_tokens": 0, + "input_tokens": 2000, + "output_tokens": 2000 + }, + "expected_cost": 500 + }, + { + "name": "no_schedule_prices_the_base_only", + "consumption_rate": 1, + "quantities": { + "input_tokens": 1000, + "output_tokens": 1000 + }, + "expected_cost": 1 + }, + { + "name": "track_event_with_three_charges", + "consumption_rate": 0.01, + "quantity_rates": { + "destination": 0.0035, + "origination": 0.004, + "transfer": 0.002 + }, + "quantity": 1, + "quantities": { + "destination": 1, + "origination": 1, + "transfer": 2 + }, + "expected_cost": 0.0215 + }, + { + "name": "decimal_quantities", + "consumption_rate": 0, + "quantity_rates": { + "gb_transferred": 0.12, + "minutes": 0.5 + }, + "quantities": { + "gb_transferred": 2.75, + "minutes": 1.5 + }, + "expected_cost": 1.08 + }, + { + "name": "rated_key_the_event_does_not_carry", + "consumption_rate": 1, + "quantity_rates": { + "origination": 2, + "transfer": 3 + }, + "quantities": { + "origination": 1 + }, + "expected_cost": 3 + }, + { + "name": "carried_key_the_schedule_does_not_rate", + "consumption_rate": 1, + "quantity_rates": { + "origination": 2 + }, + "quantities": { + "destination": 5, + "origination": 1 + }, + "expected_cost": 3 + }, + { + "name": "token_and_custom_keys_together", + "consumption_rate": 0.5, + "quantity_rates": { + "input_tokens": 0.001, + "output_tokens": 0.01, + "tool_calls": 0.2, + "web_searches": 1 + }, + "quantity": 2, + "quantities": { + "cached_input_tokens": 400, + "input_tokens": 1000, + "output_tokens": 100, + "tool_calls": 3, + "web_searches": 0.5 + }, + "expected_cost": 3.7 + }, + { + "name": "negative_quantities_price_at_zero", + "consumption_rate": 1, + "quantity_rates": { + "origination": 2 + }, + "quantities": { + "origination": -5 + }, + "expected_cost": 1 + } + ] +} diff --git a/src/SchematicHQ.Client.Test/SchematicHQ.Client.Test.Custom.props b/src/SchematicHQ.Client.Test/SchematicHQ.Client.Test.Custom.props index a7029f21..81f059c1 100644 --- a/src/SchematicHQ.Client.Test/SchematicHQ.Client.Test.Custom.props +++ b/src/SchematicHQ.Client.Test/SchematicHQ.Client.Test.Custom.props @@ -14,4 +14,7 @@ Configure additional MSBuild properties for your project in this file: + + + diff --git a/src/SchematicHQ.Client/RulesEngine/Wasm/WasmRulesEngine.cs b/src/SchematicHQ.Client/RulesEngine/Wasm/WasmRulesEngine.cs index ac77acf4..436ba6dc 100644 --- a/src/SchematicHQ.Client/RulesEngine/Wasm/WasmRulesEngine.cs +++ b/src/SchematicHQ.Client/RulesEngine/Wasm/WasmRulesEngine.cs @@ -1,5 +1,6 @@ using System; using System.IO; +using System.Linq; using System.Text; using System.Text.Json.Nodes; using Microsoft.Extensions.Logging; @@ -165,9 +166,10 @@ public CheckFlagResult CheckFlag( { envelope["user"] = Sanitize(JsonUtils.SerializeToNode(user), "user"); } - if (HasPreflight(preflight)) + var options = PreflightNode(preflight); + if (options != null) { - envelope["options"] = JsonUtils.SerializeToNode(preflight); + envelope["options"] = options; } var resultJson = CallWasm(envelope.ToJsonString()); @@ -183,16 +185,29 @@ public CheckFlagResult CheckFlag( } /// - /// Reports whether any preflight knob was supplied. With none, the - /// envelope omits options entirely and the engine uses its defaults. + /// The preflight as the engine's options envelope, or null when no knob + /// was supplied, so the envelope omits options and the engine uses its + /// defaults. Every field the body serializes goes through, not a fixed + /// list, so a knob the generated type gains (event_quantities, once the + /// API's preflight body carries it) reaches the engine without a change + /// here. A null or empty-object field counts as unset, so an empty + /// credit_cost still declares nothing. /// - private static bool HasPreflight(PreflightRequestBody? preflight) => - preflight != null - && ( - (preflight.CreditCost != null && preflight.CreditCost.Count > 0) - || preflight.Usage.HasValue - || preflight.EventUsage != null - ); + private static JsonObject? PreflightNode(PreflightRequestBody? preflight) + { + if (preflight == null || JsonUtils.SerializeToNode(preflight) is not JsonObject node) + { + return null; + } + foreach (var key in node.Select(p => p.Key).ToList()) + { + if (node[key] is null || node[key] is JsonObject { Count: 0 }) + { + node.Remove(key); + } + } + return node.Count == 0 ? null : node; + } private static CheckFlagResult ToCheckFlagResult(RulesengineCheckFlagResult r) => new()