From 058c22a47fd4931252e3ab03257a86254dc7bc0b Mon Sep 17 00:00:00 2001 From: Abiorh001 Date: Thu, 28 May 2026 12:40:48 +0100 Subject: [PATCH 1/2] Fix budget reservation safety across storage backends --- src/omniclaw/client.py | 8 +- src/omniclaw/guards/__init__.py | 17 +- src/omniclaw/guards/budget.py | 397 +++++++++++++++++++++++------- src/omniclaw/storage/file.py | 91 +++++++ src/omniclaw/storage/memory.py | 91 +++++++ src/omniclaw/storage/redis.py | 186 +++++++++++++- tests/test_guard_edge_cases.py | 57 +---- tests/test_guards.py | 91 +++++++ tests/test_payment_concurrency.py | 49 ++++ tests/test_redis_storage.py | 73 ++++++ tests/test_sdk_events.py | 29 ++- tests/test_storage_query.py | 26 ++ 12 files changed, 945 insertions(+), 170 deletions(-) create mode 100644 tests/test_redis_storage.py create mode 100644 tests/test_storage_query.py diff --git a/src/omniclaw/client.py b/src/omniclaw/client.py index 656a08c..7864a71 100644 --- a/src/omniclaw/client.py +++ b/src/omniclaw/client.py @@ -2465,8 +2465,8 @@ async def add_budget_guard( Args: wallet_id: Target wallet ID - daily_limit: Max spend per 24h - hourly_limit: Max spend per 1h + daily_limit: Max spend per calendar day bucket + hourly_limit: Max spend per calendar hour bucket total_limit: Max total spend (lifetime) name: Custom name for the guard """ @@ -2494,8 +2494,8 @@ async def add_budget_guard_for_set( Args: wallet_set_id: Target wallet set ID - daily_limit: Max spend per 24h - hourly_limit: Max spend per 1h + daily_limit: Max spend per calendar day bucket + hourly_limit: Max spend per calendar hour bucket total_limit: Max total spend (lifetime) name: Custom name for the guard """ diff --git a/src/omniclaw/guards/__init__.py b/src/omniclaw/guards/__init__.py index a37b4c9..000c6fd 100644 --- a/src/omniclaw/guards/__init__.py +++ b/src/omniclaw/guards/__init__.py @@ -10,20 +10,25 @@ Example: >>> from omniclaw.guards import BudgetGuard, SingleTxGuard, GuardChain + >>> from omniclaw.storage.memory import InMemoryStorage >>> from decimal import Decimal >>> >>> # Create guards - >>> budget = BudgetGuard(daily_limit=Decimal("100")) + >>> storage = InMemoryStorage() + >>> budget = BudgetGuard(daily_limit=Decimal("100"), storage=storage) >>> max_tx = SingleTxGuard(max_amount=Decimal("25")) >>> >>> # Combine into chain >>> chain = GuardChain([max_tx, budget]) >>> - >>> # Check payments - >>> result = await chain.check(payment_context) - >>> if result.allowed: - ... # Proceed with payment - ... budget.record_spending(payment_context.amount) + >>> # Reserve before execution, then commit or release after outcome + >>> tokens = await chain.reserve(payment_context) + >>> try: + ... # Proceed with payment execution + ... await chain.commit(tokens) + ... except Exception: + ... await chain.release(tokens) + ... raise """ from omniclaw.guards.base import ( diff --git a/src/omniclaw/guards/budget.py b/src/omniclaw/guards/budget.py index 0b8e66c..7557c53 100644 --- a/src/omniclaw/guards/budget.py +++ b/src/omniclaw/guards/budget.py @@ -6,9 +6,13 @@ from __future__ import annotations +import asyncio +import hashlib import logging +import uuid from datetime import datetime, timedelta from decimal import Decimal +from time import monotonic from typing import Any from omniclaw.events import event_emitter @@ -39,8 +43,8 @@ def __init__( At least one limit must be specified. Args: - daily_limit: Maximum spending per 24-hour rolling period - hourly_limit: Maximum spending per 1-hour rolling period + daily_limit: Maximum spending per calendar day bucket + hourly_limit: Maximum spending per calendar hour bucket total_limit: Maximum cumulative spending (no reset) name: Guard name for identification storage: Optional storage backend (if not provided, must be bound later) @@ -63,56 +67,22 @@ def name(self) -> str: return self._name async def _get_spent(self, wallet_id: str, window: timedelta | None = None) -> Decimal: - """Get spending from storage.""" + """Get committed spending from the active budget counter.""" if not self._storage: - # Fallback for tests not using storage yet, or return 0 return Decimal("0") - # We query the ledger or a specific budget collection - # For simplicity and robustness, we query the LEDGER (truth) - # Assuming we can query the associated StorageLedger via shared storage - - # Get all successful payments - # In a real optimized system, we'd use aggregation queries, but - # StorageBackend is simple CRUD. So we fetch recent records. - - # If total limit, we need everything. If rolling window, only recent. - # But since we can't query by date easily in simple KV/Doc stores without indexing, - # we might rely on a dedicated "budget_tracking" collection updated on spend. - - # STRATEGY: Use a dedicated "spending_summary" record in storage - # Key: `budget:{wallet_id}:{guard_name}` - # Data: { "total": "100.00", "history": [ {ts, amount}... ] } - # This keeps it self-contained in the guard logic. - - key = f"budget:{wallet_id}:{self.name}" - data = await self._storage.get("guard_state", key) - if not data: - return Decimal("0") - - history = data.get("history", []) - total_spent = Decimal(str(data.get("total", "0"))) - - if window is None: - return total_spent - - # Calculate window spending - cutoff = datetime.now() - window - window_spent = Decimal("0") - - for record in history: - ts = datetime.fromisoformat(record["ts"]) - if ts > cutoff: - window_spent += Decimal(str(record["amount"])) - - return window_spent + if window == timedelta(hours=1): + return await self._get_period_used(wallet_id, "hourly", include_reserved=False) + if window == timedelta(days=1): + return await self._get_period_used(wallet_id, "daily", include_reserved=False) + return await self._get_period_used(wallet_id, "total", include_reserved=False) async def get_hourly_spent(self, wallet_id: str) -> Decimal: - """Get amount spent in last hour.""" + """Get amount spent in the current calendar hour bucket.""" return await self._get_spent(wallet_id, timedelta(hours=1)) async def get_daily_spent(self, wallet_id: str) -> Decimal: - """Get amount spent in last 24 hours.""" + """Get amount spent in the current calendar day bucket.""" return await self._get_spent(wallet_id, timedelta(days=1)) async def check(self, context: PaymentContext) -> GuardResult: @@ -122,7 +92,7 @@ async def check(self, context: PaymentContext) -> GuardResult: # Check hourly limit if self._hourly_limit is not None: - hourly_spent = await self.get_hourly_spent(wallet_id) + hourly_spent = await self._get_period_used(wallet_id, "hourly") if hourly_spent + amount > self._hourly_limit: event_emitter.emit_background( "guard.budget_exceeded", wallet_id, payload={"amount": str(amount)} @@ -148,7 +118,7 @@ async def check(self, context: PaymentContext) -> GuardResult: # Check daily limit if self._daily_limit is not None: - daily_spent = await self.get_daily_spent(wallet_id) + daily_spent = await self._get_period_used(wallet_id, "daily") if daily_spent + amount > self._daily_limit: event_emitter.emit_background( "guard.budget_exceeded", wallet_id, payload={"amount": str(amount)} @@ -174,7 +144,7 @@ async def check(self, context: PaymentContext) -> GuardResult: # Check total limit if self._total_limit is not None: - total_spent = await self.get_total_spent(wallet_id) + total_spent = await self._get_period_used(wallet_id, "total") if total_spent + amount > self._total_limit: event_emitter.emit_background( "guard.budget_exceeded", wallet_id, payload={"amount": str(amount)} @@ -200,7 +170,7 @@ async def check(self, context: PaymentContext) -> GuardResult: remaining = {} if self._hourly_limit is not None: - hourly = await self.get_hourly_spent(wallet_id) + hourly = await self._get_period_used(wallet_id, "hourly") remaining["hourly"] = self._hourly_limit - hourly if remaining["hourly"] < self._hourly_limit * Decimal("0.2"): event_emitter.emit_background( @@ -209,7 +179,7 @@ async def check(self, context: PaymentContext) -> GuardResult: payload={"remaining": str(remaining["hourly"])}, ) if self._daily_limit is not None: - daily = await self.get_daily_spent(wallet_id) + daily = await self._get_period_used(wallet_id, "daily") remaining["daily"] = self._daily_limit - daily if remaining["daily"] < self._daily_limit * Decimal("0.2"): event_emitter.emit_background( @@ -245,6 +215,149 @@ def _get_period_keys(self, wallet_id: str, ts: datetime) -> dict[str, str]: return keys + def _parse_storage_decimal(self, data: Any) -> Decimal: + if data is None: + return Decimal("0") + if isinstance(data, dict): + return Decimal(str(data.get("value", "0"))) + return Decimal(str(data)) + + async def _get_period_used( + self, + wallet_id: str, + limit_type: str, + *, + include_reserved: bool = True, + ) -> Decimal: + """Read the current committed plus in-flight amount for a limit type.""" + if not self._storage: + return Decimal("0") + period_keys = self._get_period_keys(wallet_id, datetime.now()) + key_base = period_keys.get(limit_type) + if not key_base: + return Decimal("0") + + main_data = await self._storage.get("guard_state", key_base) + total = self._parse_storage_decimal(main_data) + if include_reserved: + reserved_data = await self._storage.get("guard_state", f"{key_base}:reserved") + total += self._parse_storage_decimal(reserved_data) + return total + + def _reservation_key(self, reservation_id: str) -> str: + return f"budget_reservation:{self.name}:{reservation_id}" + + def _legacy_reservation_id(self, token: str) -> str: + return hashlib.sha256(token.encode("utf-8")).hexdigest() + + def _limit_type_from_reservation_result( + self, + result: str, + period_keys: dict[str, str], + ) -> str: + suffix = result.split(":", 1)[1] if ":" in result else "" + period_values = list(period_keys.values()) + if suffix.isdigit(): + index = int(suffix) - 1 + if 0 <= index < len(period_values): + suffix = period_values[index] + + for limit_type, key in period_keys.items(): + if key == suffix: + return limit_type + return "budget" + + async def _load_reservation(self, token: str) -> tuple[str, dict[str, Any]]: + import json + + data = json.loads(token) + if data.get("v") == 2: + reservation_id = self._legacy_reservation_id(token) + key = self._reservation_key(reservation_id) + if not self._storage: + raise ValueError("Budget guard storage is not configured") + + locks = await self._acquire_period_locks({"legacy": f"legacy:{key}"}) + try: + record = await self._storage.get("guard_state", key) + if record: + return key, record + + wallet_id = str(data["w"]) + amount = str(data["a"]) + ts = datetime.fromisoformat(data["ts"]) + legacy_record = { + "id": reservation_id, + "status": "reserved", + "wallet_id": wallet_id, + "amount": amount, + "period_keys": list(self._get_period_keys(wallet_id, ts).values()), + "created_at": ts.isoformat(), + "migrated_from_token_v": 2, + } + await self._storage.save("guard_state", key, legacy_record) + return key, legacy_record + finally: + await self._release_period_locks(locks) + + if data.get("v") != 3: + raise ValueError("Unsupported budget reservation token") + + reservation_id = str(data["id"]) + key = self._reservation_key(reservation_id) + if not self._storage: + raise ValueError("Budget guard storage is not configured") + + record = await self._storage.get("guard_state", key) + if not record: + raise ValueError("Budget reservation record is missing") + return key, record + + async def _acquire_period_locks( + self, + period_keys: dict[str, str], + *, + ttl: int = 30, + timeout: float = 30.0, + ) -> list[tuple[str, str]]: + if not self._storage: + return [] + + lock_keys = [f"guard:budget:{key}" for key in sorted(period_keys.values())] + deadline = monotonic() + timeout + + while True: + acquired: list[tuple[str, str]] = [] + try: + for lock_key in lock_keys: + token = await self._storage.acquire_lock(lock_key, ttl=ttl) + if token is None: + await self._release_period_locks(acquired) + break + acquired.append((lock_key, token)) + else: + return acquired + except BaseException: + await self._release_period_locks(acquired) + raise + + if monotonic() >= deadline: + raise TimeoutError("Timed out acquiring budget reservation locks") + await asyncio.sleep(0.005) + + async def _release_period_locks(self, locks: list[tuple[str, str]]) -> None: + if not self._storage: + return + for lock_key, token in reversed(locks): + try: + released = await self._storage.release_lock(lock_key, token) + if not released: + logger.warning( + "Budget lock %s was not released; ownership may be lost", lock_key + ) + except Exception: + logger.warning("Failed to release budget lock %s", lock_key, exc_info=True) + async def reserve(self, context: PaymentContext) -> str | None: """Atomic reservation for all configured limits.""" if not self._storage: @@ -260,74 +373,137 @@ async def reserve(self, context: PaymentContext) -> str | None: return None # No limits configured reserved_keys: list[str] = [] + locks: list[tuple[str, str]] = [] + reservation_id = str(uuid.uuid4()) + reservation_key = self._reservation_key(reservation_id) + reservation_record = { + "id": reservation_id, + "status": "reserved", + "wallet_id": wallet_id, + "amount": str(amount), + "period_keys": list(period_keys.values()), + "created_at": now.isoformat(), + } + storage_create = getattr(self._storage, "create_budget_reservation", None) + + if storage_create: + period_limits = { + key_base: str(getattr(self, f"_{limit_type}_limit")) + for limit_type, key_base in period_keys.items() + } + result = await storage_create( + "guard_state", + reservation_key, + period_limits, + str(amount), + reservation_record, + ) + if result == "reserved": + import json + + return json.dumps({"v": 3, "id": reservation_id}) + if result.startswith("limit_exceeded"): + limit_type = self._limit_type_from_reservation_result(result, period_keys) + limit = getattr(self, f"_{limit_type}_limit", "unknown") + raise ValueError(f"{limit_type.capitalize()} budget limit exceeded. Limit: {limit}") + raise RuntimeError(f"Budget reservation failed closed: {result}") try: - # 2. Iterate and Reserve each + locks = await self._acquire_period_locks(period_keys) + + # 2. Check every active limit before mutating reservation state. for limit_type, key_base in period_keys.items(): key_reserved = f"{key_base}:reserved" key_main = key_base limit = getattr(self, f"_{limit_type}_limit") - # Optimistic Increment Reserved - await self._storage.atomic_add("guard_state", key_reserved, str(amount)) - reserved_keys.append(key_reserved) - - # Fetch Check main_data = await self._storage.get("guard_state", key_main) res_data = await self._storage.get("guard_state", key_reserved) - def _parse_val(d: Any) -> Decimal: - if d is None: - return Decimal("0") - if isinstance(d, dict): - return Decimal(str(d.get("value", "0"))) - return Decimal(str(d)) - - current_main = _parse_val(main_data) - current_res = _parse_val(res_data) + current_main = self._parse_storage_decimal(main_data) + current_res = self._parse_storage_decimal(res_data) - if current_main + current_res > limit: + if limit is not None and current_main + current_res + amount > limit: raise ValueError( f"{limit_type.capitalize()} budget limit exceeded. Limit: {limit}" ) + # 3. Reserve all active limits while holding storage-backed locks. + for key_base in period_keys.values(): + key_reserved = f"{key_base}:reserved" + await self._storage.atomic_add("guard_state", key_reserved, str(amount)) + reserved_keys.append(key_reserved) + + await self._storage.save("guard_state", reservation_key, reservation_record) + except Exception: # Rollback ALL reserved keys for rk in reserved_keys: await self._storage.atomic_add("guard_state", rk, str(-amount)) raise + finally: + await self._release_period_locks(locks) # Token = JSON string with context to reconstruct keys import json - token_data = {"v": 2, "w": wallet_id, "a": str(amount), "ts": now.isoformat()} + token_data = {"v": 3, "id": reservation_id} return json.dumps(token_data) async def commit(self, token: str | None) -> None: if not token or not self._storage: return - import json wallet_id = "" try: - data = json.loads(token) - if data.get("v") != 2: + reservation_key, record = await self._load_reservation(token) + status = record.get("status") + if status != "reserved": return - amount = Decimal(data["a"]) - wallet_id = data["w"] - ts = datetime.fromisoformat(data["ts"]) - - # Reconstruct Keys using ORIGINAL timestamp - # This ensures we commit to the bucket we reserved in - period_keys = self._get_period_keys(wallet_id, ts) - - for key_base in period_keys.values(): - key_reserved = f"{key_base}:reserved" - - # Move Reserved -> Main - await self._storage.atomic_add("guard_state", key_base, str(amount)) - await self._storage.atomic_add("guard_state", key_reserved, str(-amount)) + amount = Decimal(str(record["amount"])) + wallet_id = str(record["wallet_id"]) + period_key_values = [str(key) for key in record.get("period_keys", [])] + committed_at = datetime.now().isoformat() + + storage_commit = getattr(self._storage, "commit_budget_reservation", None) + if storage_commit: + result = await storage_commit( + "guard_state", + reservation_key, + period_key_values, + str(amount), + committed_at, + ) + if result in {"committed", "released"}: + return + if result == "missing": + raise ValueError("Budget reservation record is missing") + + period_keys = {str(index): key for index, key in enumerate(period_key_values)} + locks = await self._acquire_period_locks(period_keys) + + try: + current = await self._storage.get("guard_state", reservation_key) + if not current: + raise ValueError("Budget reservation record is missing") + if current.get("status") != "reserved": + return + + for key_base in period_key_values: + key_reserved = f"{key_base}:reserved" + + # Move Reserved -> Main + await self._storage.atomic_add("guard_state", key_base, str(amount)) + await self._storage.atomic_add("guard_state", key_reserved, str(-amount)) + + await self._storage.update( + "guard_state", + reservation_key, + {"status": "committed", "committed_at": committed_at}, + ) + finally: + await self._release_period_locks(locks) except Exception as exc: logger.error( @@ -342,23 +518,54 @@ async def commit(self, token: str | None) -> None: async def release(self, token: str | None) -> None: if not token or not self._storage: return - import json wallet_id = "" try: - data = json.loads(token) - if data.get("v") != 2: + reservation_key, record = await self._load_reservation(token) + status = record.get("status") + if status != "reserved": return - amount = Decimal(data["a"]) - wallet_id = data["w"] - ts = datetime.fromisoformat(data["ts"]) - - period_keys = self._get_period_keys(wallet_id, ts) - - for key_base in period_keys.values(): - key_reserved = f"{key_base}:reserved" - await self._storage.atomic_add("guard_state", key_reserved, str(-amount)) + amount = Decimal(str(record["amount"])) + wallet_id = str(record["wallet_id"]) + period_key_values = [str(key) for key in record.get("period_keys", [])] + released_at = datetime.now().isoformat() + + storage_release = getattr(self._storage, "release_budget_reservation", None) + if storage_release: + result = await storage_release( + "guard_state", + reservation_key, + period_key_values, + str(amount), + released_at, + ) + if result in {"released", "committed"}: + return + if result == "missing": + raise ValueError("Budget reservation record is missing") + + period_keys = {str(index): key for index, key in enumerate(period_key_values)} + locks = await self._acquire_period_locks(period_keys) + + try: + current = await self._storage.get("guard_state", reservation_key) + if not current: + raise ValueError("Budget reservation record is missing") + if current.get("status") != "reserved": + return + + for key_base in period_key_values: + key_reserved = f"{key_base}:reserved" + await self._storage.atomic_add("guard_state", key_reserved, str(-amount)) + + await self._storage.update( + "guard_state", + reservation_key, + {"status": "released", "released_at": released_at}, + ) + finally: + await self._release_period_locks(locks) except Exception as exc: logger.error( diff --git a/src/omniclaw/storage/file.py b/src/omniclaw/storage/file.py index bbbec4e..2bc9c83 100644 --- a/src/omniclaw/storage/file.py +++ b/src/omniclaw/storage/file.py @@ -100,6 +100,9 @@ async def query( results = [] for key, value in data.items(): + if not isinstance(value, dict): + value = {"value": value} + if filters: match = True for filter_key, filter_value in filters.items(): @@ -173,6 +176,94 @@ async def atomic_add( await self._write_collection(collection, data) return str(new_val) + async def create_budget_reservation( + self, + collection: str, + reservation_key: str, + period_limits: dict[str, str], + amount: str, + record: dict[str, Any], + ) -> str: + """Atomically create a budget reservation if every period has capacity.""" + async with self._lock: + data = await self._read_collection(collection) + amount_dec = Decimal(str(amount)) + + for period_key, limit in period_limits.items(): + main = Decimal(str(data.get(period_key, "0"))) + reserved = Decimal(str(data.get(f"{period_key}:reserved", "0"))) + if main + reserved + amount_dec > Decimal(str(limit)): + return f"limit_exceeded:{period_key}" + + data[reservation_key] = record + for period_key in period_limits: + reserved_key = f"{period_key}:reserved" + current = Decimal(str(data.get(reserved_key, "0"))) + data[reserved_key] = str(current + amount_dec) + + await self._write_collection(collection, data) + return "reserved" + + async def commit_budget_reservation( + self, + collection: str, + reservation_key: str, + period_keys: list[str], + amount: str, + committed_at: str, + ) -> str: + """Atomically commit a budget reservation and mark it single-use.""" + async with self._lock: + data = await self._read_collection(collection) + record = data.get(reservation_key) + if not record: + return "missing" + status = record.get("status") + if status != "reserved": + return str(status) + + amount_dec = Decimal(str(amount)) + for period_key in period_keys: + main = Decimal(str(data.get(period_key, "0"))) + reserved_key = f"{period_key}:reserved" + reserved = Decimal(str(data.get(reserved_key, "0"))) + data[period_key] = str(main + amount_dec) + data[reserved_key] = str(reserved - amount_dec) + + record["status"] = "committed" + record["committed_at"] = committed_at + await self._write_collection(collection, data) + return "committed" + + async def release_budget_reservation( + self, + collection: str, + reservation_key: str, + period_keys: list[str], + amount: str, + released_at: str, + ) -> str: + """Atomically release a budget reservation and mark it single-use.""" + async with self._lock: + data = await self._read_collection(collection) + record = data.get(reservation_key) + if not record: + return "missing" + status = record.get("status") + if status != "reserved": + return str(status) + + amount_dec = Decimal(str(amount)) + for period_key in period_keys: + reserved_key = f"{period_key}:reserved" + reserved = Decimal(str(data.get(reserved_key, "0"))) + data[reserved_key] = str(reserved - amount_dec) + + record["status"] = "released" + record["released_at"] = released_at + await self._write_collection(collection, data) + return "released" + async def acquire_lock( self, key: str, diff --git a/src/omniclaw/storage/memory.py b/src/omniclaw/storage/memory.py index fbb58f4..c128ce9 100644 --- a/src/omniclaw/storage/memory.py +++ b/src/omniclaw/storage/memory.py @@ -74,6 +74,9 @@ async def query( results = [] for key, data in coll.items(): + if not isinstance(data, dict): + data = {"value": data} + # Apply filters if filters: match = True @@ -161,6 +164,94 @@ async def atomic_add( coll[key] = str(new_val) return str(new_val) + async def create_budget_reservation( + self, + collection: str, + reservation_key: str, + period_limits: dict[str, str], + amount: str, + record: dict[str, Any], + ) -> str: + """Atomically create a budget reservation if every period has capacity.""" + from decimal import Decimal + + coll = self._ensure_collection(collection) + amount_dec = Decimal(str(amount)) + + for period_key, limit in period_limits.items(): + main = Decimal(str(coll.get(period_key, "0"))) + reserved = Decimal(str(coll.get(f"{period_key}:reserved", "0"))) + if main + reserved + amount_dec > Decimal(str(limit)): + return f"limit_exceeded:{period_key}" + + coll[reservation_key] = deepcopy(record) + for period_key in period_limits: + reserved_key = f"{period_key}:reserved" + current = Decimal(str(coll.get(reserved_key, "0"))) + coll[reserved_key] = str(current + amount_dec) + + return "reserved" + + async def commit_budget_reservation( + self, + collection: str, + reservation_key: str, + period_keys: list[str], + amount: str, + committed_at: str, + ) -> str: + """Atomically commit a budget reservation and mark it single-use.""" + from decimal import Decimal + + coll = self._ensure_collection(collection) + record = coll.get(reservation_key) + if not record: + return "missing" + status = record.get("status") + if status != "reserved": + return str(status) + + amount_dec = Decimal(str(amount)) + for period_key in period_keys: + main = Decimal(str(coll.get(period_key, "0"))) + reserved_key = f"{period_key}:reserved" + reserved = Decimal(str(coll.get(reserved_key, "0"))) + coll[period_key] = str(main + amount_dec) + coll[reserved_key] = str(reserved - amount_dec) + + record["status"] = "committed" + record["committed_at"] = committed_at + return "committed" + + async def release_budget_reservation( + self, + collection: str, + reservation_key: str, + period_keys: list[str], + amount: str, + released_at: str, + ) -> str: + """Atomically release a budget reservation and mark it single-use.""" + from decimal import Decimal + + coll = self._ensure_collection(collection) + record = coll.get(reservation_key) + if not record: + return "missing" + status = record.get("status") + if status != "reserved": + return str(status) + + amount_dec = Decimal(str(amount)) + for period_key in period_keys: + reserved_key = f"{period_key}:reserved" + reserved = Decimal(str(coll.get(reserved_key, "0"))) + coll[reserved_key] = str(reserved - amount_dec) + + record["status"] = "released" + record["released_at"] = released_at + return "released" + async def acquire_lock( self, key: str, diff --git a/src/omniclaw/storage/redis.py b/src/omniclaw/storage/redis.py index 918f43a..6fb9833 100644 --- a/src/omniclaw/storage/redis.py +++ b/src/omniclaw/storage/redis.py @@ -102,7 +102,10 @@ async def get( return None try: - return json.loads(data) + parsed = json.loads(data) + if isinstance(parsed, dict): + return parsed + return {"value": parsed} except json.JSONDecodeError: # Fallback for keys created via atomic_add (raw strings) return {"value": data} @@ -170,6 +173,94 @@ async def atomic_add( end """ + _BUDGET_RESERVE_SCRIPT = """ + if redis.call("get", KEYS[1]) then + return "exists" + end + + local amount = tonumber(ARGV[1]) + local amount_units = tonumber(ARGV[4]) + local count = tonumber(ARGV[3]) + local scale = tonumber(ARGV[5]) + if not amount or not amount_units or not count or not scale then + return "invalid_input" + end + + for i = 1, count do + local main_value = redis.call("get", KEYS[1 + i]) or "0" + local reserved_value = redis.call("get", KEYS[1 + count + i]) or "0" + local limit_units = tonumber(ARGV[5 + i]) + local main = tonumber(main_value) + local reserved = tonumber(reserved_value) + + if not main or not reserved or not limit_units then + return "counter_corrupt" + end + + local used_units = math.floor((main + reserved) * scale + 0.5) + if used_units + amount_units > limit_units then + return "limit_exceeded:" .. tostring(i) + end + end + + redis.call("set", KEYS[1], ARGV[2]) + + for i = 1, count do + redis.call("INCRBYFLOAT", KEYS[1 + count + i], ARGV[1]) + end + + return "reserved" + """ + + _BUDGET_COMMIT_SCRIPT = """ + local record = redis.call("get", KEYS[1]) + if not record then + return "missing" + end + + local data = cjson.decode(record) + if data["status"] ~= "reserved" then + return data["status"] + end + + local amount = ARGV[1] + local count = tonumber(ARGV[3]) + + for i = 1, count do + redis.call("INCRBYFLOAT", KEYS[1 + i], amount) + redis.call("INCRBYFLOAT", KEYS[1 + count + i], "-" .. amount) + end + + data["status"] = "committed" + data["committed_at"] = ARGV[2] + redis.call("set", KEYS[1], cjson.encode(data)) + return "committed" + """ + + _BUDGET_RELEASE_SCRIPT = """ + local record = redis.call("get", KEYS[1]) + if not record then + return "missing" + end + + local data = cjson.decode(record) + if data["status"] ~= "reserved" then + return data["status"] + end + + local amount = ARGV[1] + local count = tonumber(ARGV[3]) + + for i = 1, count do + redis.call("INCRBYFLOAT", KEYS[1 + i], "-" .. amount) + end + + data["status"] = "released" + data["released_at"] = ARGV[2] + redis.call("set", KEYS[1], cjson.encode(data)) + return "released" + """ + async def acquire_lock( self, key: str, @@ -232,6 +323,97 @@ async def refresh_lock( result = await client.eval(self._REFRESH_LOCK_SCRIPT, 1, redis_key, token, ttl) return int(result) > 0 + async def create_budget_reservation( + self, + collection: str, + reservation_key: str, + period_limits: dict[str, str], + amount: str, + record: dict[str, Any], + ) -> str: + """Atomically create a budget reservation if every period has capacity.""" + client = self._get_client() + period_keys = list(period_limits.keys()) + main_keys = [self._make_key(collection, key) for key in period_keys] + reserved_keys = [self._make_key(collection, f"{key}:reserved") for key in period_keys] + keys = [self._make_key(collection, reservation_key), *main_keys, *reserved_keys] + amount_str = str(Decimal(str(amount))) + scale = Decimal("1000000") + amount_units = str(int((Decimal(amount_str) * scale).to_integral_exact())) + result = await client.eval( + self._BUDGET_RESERVE_SCRIPT, + len(keys), + *keys, + amount_str, + json.dumps(record), + str(len(period_keys)), + amount_units, + str(int(scale)), + *[ + str(int((Decimal(str(limit)) * scale).to_integral_exact())) + for limit in period_limits.values() + ], + ) + + if result == "reserved": + index_key = f"{self._prefix}:{collection}:_index" + await client.sadd( + index_key, + reservation_key, + *period_keys, + *[f"{key}:reserved" for key in period_keys], + ) + return str(result) + + async def commit_budget_reservation( + self, + collection: str, + reservation_key: str, + period_keys: list[str], + amount: str, + committed_at: str, + ) -> str: + """Atomically commit a budget reservation and mark it single-use.""" + client = self._get_client() + main_keys = [self._make_key(collection, key) for key in period_keys] + reserved_keys = [self._make_key(collection, f"{key}:reserved") for key in period_keys] + keys = [self._make_key(collection, reservation_key), *main_keys, *reserved_keys] + result = await client.eval( + self._BUDGET_COMMIT_SCRIPT, + len(keys), + *keys, + str(Decimal(str(amount))), + committed_at, + str(len(period_keys)), + ) + + if result == "committed": + index_key = f"{self._prefix}:{collection}:_index" + await client.sadd(index_key, *period_keys, *[f"{key}:reserved" for key in period_keys]) + return str(result) + + async def release_budget_reservation( + self, + collection: str, + reservation_key: str, + period_keys: list[str], + amount: str, + released_at: str, + ) -> str: + """Atomically release a budget reservation and mark it single-use.""" + client = self._get_client() + reserved_keys = [self._make_key(collection, f"{key}:reserved") for key in period_keys] + keys = [self._make_key(collection, reservation_key), *reserved_keys] + result = await client.eval( + self._BUDGET_RELEASE_SCRIPT, + len(keys), + *keys, + str(Decimal(str(amount))), + released_at, + str(len(period_keys)), + ) + return str(result) + async def query( self, collection: str, @@ -251,6 +433,8 @@ async def query( data = await self.get(collection, key) if data is None: continue + if not isinstance(data, dict): + data = {"value": data} # Apply filters if filters: diff --git a/tests/test_guard_edge_cases.py b/tests/test_guard_edge_cases.py index 6aa8f0e..497b96f 100644 --- a/tests/test_guard_edge_cases.py +++ b/tests/test_guard_edge_cases.py @@ -3,13 +3,13 @@ """ from decimal import Decimal -from unittest.mock import AsyncMock, MagicMock import pytest from omniclaw.guards.base import PaymentContext from omniclaw.guards.budget import BudgetGuard from omniclaw.guards.single_tx import SingleTxGuard +from omniclaw.storage.memory import InMemoryStorage @pytest.fixture @@ -24,18 +24,7 @@ def mock_context(): @pytest.mark.asyncio async def test_budget_exact_limit(mock_context): """Test payment exactly equal to the budget limit.""" - guard = BudgetGuard(daily_limit=Decimal("100.00"), name="budget") - guard._storage = MagicMock() - - # Mock get to return 100.00 for reserved key check - async def mock_get(collection, key): - if key.endswith(":reserved"): - return "100.00" - return None - - guard._storage.get = AsyncMock(side_effect=mock_get) - guard._storage.atomic_add = AsyncMock(return_value=Decimal("100.00")) - + guard = BudgetGuard(daily_limit=Decimal("100.00"), name="budget", storage=InMemoryStorage()) mock_context.amount = Decimal("100.00") # Reserve should succeed (100 <= 100) @@ -46,28 +35,16 @@ async def mock_get(collection, key): @pytest.mark.asyncio async def test_budget_exceeds_by_smallest_unit(mock_context): """Test payment exceeding budget by 0.01.""" - guard = BudgetGuard(daily_limit=Decimal("100.00"), name="budget") - guard._storage = MagicMock() - - async def mock_get(collection, key): - if key.endswith(":reserved"): - return "100.01" # Simulating the result of atomic_add - return None - - guard._storage.get = AsyncMock(side_effect=mock_get) - guard._storage.atomic_add = AsyncMock(return_value=Decimal("100.01")) - + storage = InMemoryStorage() + guard = BudgetGuard(daily_limit=Decimal("100.00"), name="budget", storage=storage) mock_context.amount = Decimal("100.01") # Reserve should fail by raising ValueError with pytest.raises(ValueError, match="budget limit exceeded"): await guard.reserve(mock_context) - # Should revert (atomic_add negative amount) - assert guard._storage.atomic_add.call_count == 2 - args_2 = guard._storage.atomic_add.call_args_list[1] - # Check that second call added negative amount - assert args_2[0][2] == str(Decimal("-100.01")) + guard_state = await storage.query("guard_state") + assert not [entry for entry in guard_state if entry["_key"].endswith(":reserved")] @pytest.mark.asyncio @@ -87,16 +64,7 @@ async def test_single_tx_exact_limit(mock_context): @pytest.mark.asyncio async def test_negative_amount_handling(mock_context): """Test guards handling negative amounts.""" - guard = BudgetGuard(daily_limit=Decimal("100.00"), name="budget") - guard._storage = MagicMock() - - async def mock_get(collection, key): - if key.endswith(":reserved"): - return "-10.00" - return None - - guard._storage.get = AsyncMock(side_effect=mock_get) - guard._storage.atomic_add = AsyncMock(return_value=Decimal("-10.00")) + guard = BudgetGuard(daily_limit=Decimal("100.00"), name="budget", storage=InMemoryStorage()) mock_context.amount = Decimal("-10.00") @@ -107,16 +75,7 @@ async def mock_get(collection, key): @pytest.mark.asyncio async def test_zero_amount_budget(mock_context): """Test zero amount payment impacting budget.""" - guard = BudgetGuard(daily_limit=Decimal("100.00"), name="budget") - guard._storage = MagicMock() - - async def mock_get(collection, key): - if key.endswith(":reserved"): - return "50.00" # Assuming some previous usage? or just 0 + 0 - return None - - guard._storage.get = AsyncMock(side_effect=mock_get) - guard._storage.atomic_add = AsyncMock(return_value=Decimal("50.00")) + guard = BudgetGuard(daily_limit=Decimal("100.00"), name="budget", storage=InMemoryStorage()) mock_context.amount = Decimal("0") diff --git a/tests/test_guards.py b/tests/test_guards.py index 1087ae4..d1daaf3 100644 --- a/tests/test_guards.py +++ b/tests/test_guards.py @@ -103,6 +103,97 @@ async def test_tracks_spending(self, payment_context): with pytest.raises(ValueError, match="Daily budget limit exceeded"): await guard.reserve(payment_context) + @pytest.mark.asyncio + async def test_check_counts_reserved_capacity(self, payment_context): + storage = InMemoryStorage() + guard = BudgetGuard(daily_limit=Decimal("15.00"), storage=storage) + + token = await guard.reserve(payment_context) + assert token is not None + + result = await guard.check(payment_context) + assert result.allowed is False + assert "daily" in result.reason.lower() + + @pytest.mark.asyncio + async def test_commit_is_idempotent(self): + storage = InMemoryStorage() + guard = BudgetGuard(total_limit=Decimal("100.00"), storage=storage) + context = PaymentContext(wallet_id="wallet-123", recipient="0x123", amount=Decimal("10.00")) + + token = await guard.reserve(context) + await guard.commit(token) + await guard.commit(token) + + total = await storage.get("guard_state", "budget:wallet-123:budget:total") + reserved = await storage.get("guard_state", "budget:wallet-123:budget:total:reserved") + assert total == "10.00" + assert reserved == "0.00" + + @pytest.mark.asyncio + async def test_release_is_idempotent(self): + storage = InMemoryStorage() + guard = BudgetGuard(total_limit=Decimal("100.00"), storage=storage) + context = PaymentContext(wallet_id="wallet-123", recipient="0x123", amount=Decimal("10.00")) + + token = await guard.reserve(context) + await guard.release(token) + await guard.release(token) + + total = await storage.get("guard_state", "budget:wallet-123:budget:total") + reserved = await storage.get("guard_state", "budget:wallet-123:budget:total:reserved") + assert total is None + assert reserved == "0.00" + + @pytest.mark.asyncio + async def test_commit_after_release_does_not_spend(self): + storage = InMemoryStorage() + guard = BudgetGuard(total_limit=Decimal("100.00"), storage=storage) + context = PaymentContext(wallet_id="wallet-123", recipient="0x123", amount=Decimal("10.00")) + + token = await guard.reserve(context) + await guard.release(token) + await guard.commit(token) + + total = await storage.get("guard_state", "budget:wallet-123:budget:total") + reserved = await storage.get("guard_state", "budget:wallet-123:budget:total:reserved") + assert total is None + assert reserved == "0.00" + + @pytest.mark.asyncio + async def test_legacy_reservation_token_can_commit_once(self): + import json + + storage = InMemoryStorage() + guard = BudgetGuard(total_limit=Decimal("100.00"), storage=storage) + await storage.atomic_add("guard_state", "budget:wallet-123:budget:total:reserved", "10.00") + token = json.dumps({"v": 2, "w": "wallet-123", "a": "10.00", "ts": "2026-05-28T10:00:00"}) + + await guard.commit(token) + await guard.commit(token) + + total = await storage.get("guard_state", "budget:wallet-123:budget:total") + reserved = await storage.get("guard_state", "budget:wallet-123:budget:total:reserved") + assert total == "10.00" + assert reserved == "0.00" + + @pytest.mark.asyncio + async def test_legacy_reservation_token_can_release_once(self): + import json + + storage = InMemoryStorage() + guard = BudgetGuard(total_limit=Decimal("100.00"), storage=storage) + await storage.atomic_add("guard_state", "budget:wallet-123:budget:total:reserved", "10.00") + token = json.dumps({"v": 2, "w": "wallet-123", "a": "10.00", "ts": "2026-05-28T10:00:00"}) + + await guard.release(token) + await guard.release(token) + + total = await storage.get("guard_state", "budget:wallet-123:budget:total") + reserved = await storage.get("guard_state", "budget:wallet-123:budget:total:reserved") + assert total is None + assert reserved == "0.00" + @pytest.mark.asyncio async def test_hourly_limit(self, payment_context): guard = BudgetGuard(hourly_limit=Decimal("5.00"), storage=InMemoryStorage()) diff --git a/tests/test_payment_concurrency.py b/tests/test_payment_concurrency.py index 36d74d1..8726145 100644 --- a/tests/test_payment_concurrency.py +++ b/tests/test_payment_concurrency.py @@ -1,15 +1,29 @@ import asyncio from decimal import Decimal +from typing import Any from unittest.mock import AsyncMock, MagicMock import pytest from omniclaw.client import OmniClaw from omniclaw.core.types import Network, PaymentResult, PaymentStatus +from omniclaw.guards.base import PaymentContext from omniclaw.guards.budget import BudgetGuard from omniclaw.storage.memory import InMemoryStorage +class YieldingInMemoryStorage(InMemoryStorage): + """In-memory storage that yields around operations to expose async races.""" + + async def atomic_add(self, collection: str, key: str, amount: str) -> str: + await asyncio.sleep(0) + return await super().atomic_add(collection, key, amount) + + async def get(self, collection: str, key: str) -> dict[str, Any] | None: + await asyncio.sleep(0) + return await super().get(collection, key) + + @pytest.fixture def client_with_storage(): """Create a client with real in-memory storage for concurrency testing.""" @@ -92,3 +106,38 @@ async def make_payment(): assert success_count <= 16, f"Budget exceeded! {success_count} payments succeeded." assert success_count + failed_count + exception_count == 20 + + +@pytest.mark.asyncio +async def test_budget_guard_concurrent_reservations_do_not_under_authorize_with_async_storage(): + """Concurrent reservations should admit capacity instead of all rolling back.""" + storage = YieldingInMemoryStorage() + guard = BudgetGuard(total_limit=Decimal("100.00"), name="concurrent_budget", storage=storage) + + async def reserve_and_commit(index: int) -> str: + context = PaymentContext( + wallet_id="wallet-concurrent", + recipient="0x742d35Cc6634C0532925a3b844Bc9e7595f5e4a0", + amount=Decimal("6.00"), + purpose=f"concurrent-{index}", + ) + try: + token = await guard.reserve(context) + except ValueError: + return "rejected" + + await asyncio.sleep(0.01) + await guard.commit(token) + return "success" + + results = await asyncio.gather(*(reserve_and_commit(i) for i in range(20))) + + assert results.count("success") == 16 + assert results.count("rejected") == 4 + + total = await storage.get("guard_state", "budget:wallet-concurrent:concurrent_budget:total") + reserved = await storage.get( + "guard_state", "budget:wallet-concurrent:concurrent_budget:total:reserved" + ) + assert total == "96.00" + assert reserved == "0.00" diff --git a/tests/test_redis_storage.py b/tests/test_redis_storage.py new file mode 100644 index 0000000..f168423 --- /dev/null +++ b/tests/test_redis_storage.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +import uuid +from decimal import Decimal + +import pytest + +from omniclaw.guards.base import PaymentContext +from omniclaw.guards.budget import BudgetGuard +from omniclaw.storage.redis import RedisStorage + + +async def _redis_or_skip() -> RedisStorage: + storage = RedisStorage(prefix=f"omniclaw-test-{uuid.uuid4()}") + if not await storage.health_check(): + pytest.skip("Redis server is not available") + return storage + + +@pytest.mark.asyncio +async def test_redis_numeric_counters_are_queryable(): + storage = await _redis_or_skip() + try: + await storage.atomic_add("guard_state", "budget:w1:budget:total", "10") + + value = await storage.get("guard_state", "budget:w1:budget:total") + rows = await storage.query("guard_state") + + assert value == {"value": 10} + assert any(row["_key"] == "budget:w1:budget:total" for row in rows) + finally: + await storage.clear("guard_state") + await storage.close() + + +@pytest.mark.asyncio +async def test_redis_budget_reservation_lifecycle_is_single_use(): + storage = await _redis_or_skip() + try: + guard = BudgetGuard(total_limit=Decimal("100.00"), storage=storage) + context = PaymentContext(wallet_id="wallet-redis", recipient="0x123", amount=Decimal("10.00")) + + token = await guard.reserve(context) + await guard.commit(token) + await guard.commit(token) + + total = await storage.get("guard_state", "budget:wallet-redis:budget:total") + reserved = await storage.get("guard_state", "budget:wallet-redis:budget:total:reserved") + assert total == {"value": 10} + assert reserved == {"value": 0} + finally: + await storage.clear("guard_state") + await storage.close() + + +@pytest.mark.asyncio +async def test_redis_budget_reservation_allows_exact_decimal_boundary(): + storage = await _redis_or_skip() + try: + guard = BudgetGuard(total_limit=Decimal("0.30"), storage=storage) + first = PaymentContext(wallet_id="wallet-redis", recipient="0x123", amount=Decimal("0.10")) + second = PaymentContext(wallet_id="wallet-redis", recipient="0x123", amount=Decimal("0.20")) + + token = await guard.reserve(first) + await guard.commit(token) + token = await guard.reserve(second) + await guard.commit(token) + + total = await storage.get("guard_state", "budget:wallet-redis:budget:total") + assert total == {"value": 0.3} + finally: + await storage.clear("guard_state") + await storage.close() diff --git a/tests/test_sdk_events.py b/tests/test_sdk_events.py index 9992228..6006610 100644 --- a/tests/test_sdk_events.py +++ b/tests/test_sdk_events.py @@ -9,7 +9,6 @@ from __future__ import annotations -from datetime import datetime from decimal import Decimal from unittest.mock import AsyncMock, patch @@ -55,15 +54,15 @@ async def test_check_pass_emits_guard_evaluated(self, mock_emitter): async def test_check_exceed_emits_budget_exceeded(self, mock_emitter): from omniclaw.guards.base import PaymentContext from omniclaw.guards.budget import BudgetGuard + from omniclaw.storage.memory import InMemoryStorage - storage = AsyncMock() - # Mock get to return high spend data - storage.get.return_value = { - "total": "990", - "history": [{"ts": datetime.now().isoformat(), "amount": "990"}], - } - + storage = InMemoryStorage() guard = BudgetGuard(daily_limit=Decimal("1000"), storage=storage) + spent_context = PaymentContext( + wallet_id="w-1", recipient="0xabc", amount=Decimal("990"), purpose="seed" + ) + token = await guard.reserve(spent_context) + await guard.commit(token) context = PaymentContext( wallet_id="w-1", recipient="0xabc", amount=Decimal("50"), purpose="test" @@ -80,15 +79,15 @@ async def test_check_approaching_emits_warning(self, mock_emitter): """Budget > 80% but still allowed should emit approaching warning.""" from omniclaw.guards.base import PaymentContext from omniclaw.guards.budget import BudgetGuard + from omniclaw.storage.memory import InMemoryStorage - storage = AsyncMock() - # 85% consumed (850 of 1000) - storage.get.return_value = { - "total": "850", - "history": [{"ts": datetime.now().isoformat(), "amount": "850"}], - } - + storage = InMemoryStorage() guard = BudgetGuard(daily_limit=Decimal("1000"), storage=storage) + spent_context = PaymentContext( + wallet_id="w-1", recipient="0xabc", amount=Decimal("850"), purpose="seed" + ) + token = await guard.reserve(spent_context) + await guard.commit(token) context = PaymentContext( wallet_id="w-1", recipient="0xabc", amount=Decimal("10"), purpose="test" diff --git a/tests/test_storage_query.py b/tests/test_storage_query.py new file mode 100644 index 0000000..9c097ea --- /dev/null +++ b/tests/test_storage_query.py @@ -0,0 +1,26 @@ +from __future__ import annotations + +import pytest + +from omniclaw.storage.file import FileStorage +from omniclaw.storage.memory import InMemoryStorage + + +@pytest.mark.asyncio +async def test_memory_query_handles_scalar_atomic_counters(): + storage = InMemoryStorage() + await storage.atomic_add("guard_state", "counter", "10") + + rows = await storage.query("guard_state") + + assert rows == [{"value": "10", "_key": "counter"}] + + +@pytest.mark.asyncio +async def test_file_query_handles_scalar_atomic_counters(tmp_path): + storage = FileStorage(base_dir=tmp_path) + await storage.atomic_add("guard_state", "counter", "10") + + rows = await storage.query("guard_state") + + assert rows == [{"value": "10", "_key": "counter"}] From 3b0473ff53d98d7afd4891e9c1c3135645fe26bd Mon Sep 17 00:00:00 2001 From: Abiorh001 Date: Thu, 28 May 2026 12:51:43 +0100 Subject: [PATCH 2/2] Remove old budget reservation compatibility --- src/omniclaw/guards/budget.py | 33 --------------------------------- src/omniclaw/storage/redis.py | 2 +- tests/test_guards.py | 34 ---------------------------------- tests/test_redis_storage.py | 4 +++- 4 files changed, 4 insertions(+), 69 deletions(-) diff --git a/src/omniclaw/guards/budget.py b/src/omniclaw/guards/budget.py index 7557c53..02a96a1 100644 --- a/src/omniclaw/guards/budget.py +++ b/src/omniclaw/guards/budget.py @@ -7,7 +7,6 @@ from __future__ import annotations import asyncio -import hashlib import logging import uuid from datetime import datetime, timedelta @@ -247,9 +246,6 @@ async def _get_period_used( def _reservation_key(self, reservation_id: str) -> str: return f"budget_reservation:{self.name}:{reservation_id}" - def _legacy_reservation_id(self, token: str) -> str: - return hashlib.sha256(token.encode("utf-8")).hexdigest() - def _limit_type_from_reservation_result( self, result: str, @@ -271,35 +267,6 @@ async def _load_reservation(self, token: str) -> tuple[str, dict[str, Any]]: import json data = json.loads(token) - if data.get("v") == 2: - reservation_id = self._legacy_reservation_id(token) - key = self._reservation_key(reservation_id) - if not self._storage: - raise ValueError("Budget guard storage is not configured") - - locks = await self._acquire_period_locks({"legacy": f"legacy:{key}"}) - try: - record = await self._storage.get("guard_state", key) - if record: - return key, record - - wallet_id = str(data["w"]) - amount = str(data["a"]) - ts = datetime.fromisoformat(data["ts"]) - legacy_record = { - "id": reservation_id, - "status": "reserved", - "wallet_id": wallet_id, - "amount": amount, - "period_keys": list(self._get_period_keys(wallet_id, ts).values()), - "created_at": ts.isoformat(), - "migrated_from_token_v": 2, - } - await self._storage.save("guard_state", key, legacy_record) - return key, legacy_record - finally: - await self._release_period_locks(locks) - if data.get("v") != 3: raise ValueError("Unsupported budget reservation token") diff --git a/src/omniclaw/storage/redis.py b/src/omniclaw/storage/redis.py index 6fb9833..8cfdb1f 100644 --- a/src/omniclaw/storage/redis.py +++ b/src/omniclaw/storage/redis.py @@ -307,7 +307,7 @@ async def release_lock( result = await client.eval(self._RELEASE_LOCK_SCRIPT, 1, redis_key, token) return int(result) > 0 else: - # Fallback: simple delete (legacy callers) + # Tokenless release is unsafe for shared locks but kept for the base API contract. result = await client.delete(redis_key) return result > 0 diff --git a/tests/test_guards.py b/tests/test_guards.py index d1daaf3..71bc158 100644 --- a/tests/test_guards.py +++ b/tests/test_guards.py @@ -160,40 +160,6 @@ async def test_commit_after_release_does_not_spend(self): assert total is None assert reserved == "0.00" - @pytest.mark.asyncio - async def test_legacy_reservation_token_can_commit_once(self): - import json - - storage = InMemoryStorage() - guard = BudgetGuard(total_limit=Decimal("100.00"), storage=storage) - await storage.atomic_add("guard_state", "budget:wallet-123:budget:total:reserved", "10.00") - token = json.dumps({"v": 2, "w": "wallet-123", "a": "10.00", "ts": "2026-05-28T10:00:00"}) - - await guard.commit(token) - await guard.commit(token) - - total = await storage.get("guard_state", "budget:wallet-123:budget:total") - reserved = await storage.get("guard_state", "budget:wallet-123:budget:total:reserved") - assert total == "10.00" - assert reserved == "0.00" - - @pytest.mark.asyncio - async def test_legacy_reservation_token_can_release_once(self): - import json - - storage = InMemoryStorage() - guard = BudgetGuard(total_limit=Decimal("100.00"), storage=storage) - await storage.atomic_add("guard_state", "budget:wallet-123:budget:total:reserved", "10.00") - token = json.dumps({"v": 2, "w": "wallet-123", "a": "10.00", "ts": "2026-05-28T10:00:00"}) - - await guard.release(token) - await guard.release(token) - - total = await storage.get("guard_state", "budget:wallet-123:budget:total") - reserved = await storage.get("guard_state", "budget:wallet-123:budget:total:reserved") - assert total is None - assert reserved == "0.00" - @pytest.mark.asyncio async def test_hourly_limit(self, payment_context): guard = BudgetGuard(hourly_limit=Decimal("5.00"), storage=InMemoryStorage()) diff --git a/tests/test_redis_storage.py b/tests/test_redis_storage.py index f168423..1641e02 100644 --- a/tests/test_redis_storage.py +++ b/tests/test_redis_storage.py @@ -38,7 +38,9 @@ async def test_redis_budget_reservation_lifecycle_is_single_use(): storage = await _redis_or_skip() try: guard = BudgetGuard(total_limit=Decimal("100.00"), storage=storage) - context = PaymentContext(wallet_id="wallet-redis", recipient="0x123", amount=Decimal("10.00")) + context = PaymentContext( + wallet_id="wallet-redis", recipient="0x123", amount=Decimal("10.00") + ) token = await guard.reserve(context) await guard.commit(token)