inject host time into datastream wasm rules engine (SCHY-471)

src/schematic/datastream/rules_engine.py

@@ -3,6 +3,7 @@
 import json
 import logging
 import re
+import time
 from pathlib import Path
 from typing import Any, Optional
 
@@ -59,6 +60,7 @@ def __init__(self, *, wasm_path: Optional[str] = None) -> None:
         self._alloc_fn: Any = None
         self._dealloc_fn: Any = None
         self._check_flag_fn: Any = None
+        self._set_time_fn: Any = None
         self._get_result_json_fn: Any = None
         self._get_result_json_length_fn: Any = None
         self._get_version_key_fn: Any = None
@@ -99,6 +101,10 @@ async def initialize(self) -> None:
         self._alloc_fn = exports.get("alloc")
         self._dealloc_fn = exports.get("dealloc")
         self._check_flag_fn = exports.get("checkFlagCombined")
+        # Optional: feed the wasm the current wall-clock time before each check.
+        # The raw wasm has no clock under wasmtime (SCHY-471); without it,
+        # metric-period reset timestamps are omitted. Absent on older wasm.
+        self._set_time_fn = exports.get("setCurrentTimeMillis")
         self._get_result_json_fn = exports.get("getResultJson")
         self._get_result_json_length_fn = exports.get("getResultJsonLength")
         self._get_version_key_fn = exports.get("get_version_key_wasm")
@@ -180,6 +186,11 @@ def _call_wasm(self, input_json: str) -> str:
         try:
             self._memory.write(self._store, data, ptr)
 
+            # Supply the host's current time so the engine can compute
+            # metric-period reset timestamps (SCHY-471). No-op on older wasm.
+            if self._set_time_fn is not None:
+                self._set_time_fn(self._store, int(time.time() * 1000))
+
             result_len = self._check_flag_fn(self._store, ptr, length)
             if result_len < 0:
                 raise RuntimeError("WASM checkFlagCombined returned error code")

tests/datastream/test_rules_engine.py

@@ -3,7 +3,12 @@
 import pytest
 
 from schematic.datastream.rules_engine import RulesEngineClient
-from schematic.types import RulesengineCheckFlagResult, RulesengineFlag, RulesengineRule
+from schematic.types import (
+    RulesengineCheckFlagResult,
+    RulesengineCompany,
+    RulesengineFlag,
+    RulesengineRule,
+)
 
 # Skip all tests if wasmtime is not installed
 wasmtime = pytest.importorskip("wasmtime", reason="wasmtime not installed")
@@ -130,6 +135,110 @@ async def test_returns_default_for_empty_rules(self, engine: RulesEngineClient)
         assert result.flag_key == "empty-rules"
 
 
+class TestRulesEngineClockRegression:
+    """Regression for SCHY-471.
+
+    A company-override entitlement rule whose metric condition uses a
+    calendar/billing metric period drives the engine into the
+    metric-period-reset code path, which calls ``Utc::now()``. On the raw
+    ``wasm32-unknown-unknown`` build that has no clock, so the wasm used to trap
+    (``wasm unreachable``); the SDK surfaced it as ``WASM check_flag failed`` and
+    the caller wrongly saw the flag's ``default_value`` (false) for a company
+    that is legitimately entitled. The host now injects the current time via
+    ``setCurrentTimeMillis`` so this path evaluates cleanly.
+    """
+
+    def _entitled_company(self) -> RulesengineCompany:
+        from schematic.types import (
+            RulesengineCompanyMetric,
+            RulesengineCondition,
+            RulesengineRule,
+        )
+
+        company_id = "co_entitled"
+        company_condition = RulesengineCondition(
+            id="cond_company",
+            account_id="acc_1",
+            environment_id="env_1",
+            condition_type="company",
+            operator="eq",
+            resource_ids=[company_id],
+            trait_value="",
+        )
+        # Usage 40 < allocation 100 -> override grants the feature.
+        metric_condition = RulesengineCondition(
+            id="cond_metric",
+            account_id="acc_1",
+            environment_id="env_1",
+            condition_type="metric",
+            operator="lt",
+            resource_ids=[],
+            event_subtype="api-calls",
+            metric_value=100,
+            metric_period="current_month",
+            metric_period_month_reset="billing_cycle",
+            trait_value="100",
+        )
+        override_rule = RulesengineRule(
+            id="rule_override",
+            flag_id="flag1",
+            account_id="acc_1",
+            environment_id="env_1",
+            name="Company Override",
+            rule_type="company_override",
+            value=True,
+            priority=0,
+            conditions=[company_condition, metric_condition],
+            condition_groups=[],
+        )
+        metric = RulesengineCompanyMetric(
+            account_id="acc_1",
+            environment_id="env_1",
+            company_id=company_id,
+            event_subtype="api-calls",
+            period="current_month",
+            month_reset="billing_cycle",
+            value=40,
+            created_at="2023-01-01T00:00:00Z",
+        )
+        return RulesengineCompany(
+            id=company_id,
+            account_id="acc_1",
+            environment_id="env_1",
+            keys={"id": company_id},
+            traits=[],
+            metrics=[metric],
+            rules=[override_rule],
+            entitlements=[],
+            billing_product_ids=[],
+            credit_balances={},
+            plan_ids=[],
+            plan_version_ids=[],
+        )
+
+    async def test_billing_metric_override_evaluates_without_trapping(self) -> None:
+        engine = RulesEngineClient()
+        await engine.initialize()
+
+        flag = _make_flag(id="flag1", key="mcp-access", default_value=False)
+        # Must not raise (the bug surfaced as RuntimeError "WASM flag check failed").
+        result = engine.check_flag(flag, self._entitled_company())
+
+        # The override grants the feature; a fallback to default_value would be False.
+        assert result.value is True
+        assert result.reason and "override" in result.reason.lower()
+
+    async def test_billing_metric_override_populates_reset_at(self) -> None:
+        engine = RulesEngineClient()
+        await engine.initialize()
+
+        flag = _make_flag(id="flag1", key="mcp-access", default_value=False)
+        result = engine.check_flag(flag, self._entitled_company())
+
+        # setCurrentTimeMillis lets the engine compute the next reset boundary.
+        assert result.feature_usage_reset_at is not None
+
+
 class TestRulesEngineFileNotFound:
     async def test_missing_wasm_raises(self) -> None:
         engine = RulesEngineClient(wasm_path="/nonexistent/rulesengine.wasm")