Merge various work from Orchid Labs master branch.

gai-backend/account.py

@@ -0,0 +1,141 @@
+from web3 import Web3
+from decimal import Decimal
+import secrets
+from eth_account.messages import encode_defunct
+from lottery import Lottery
+from typing import Optional, Dict, Tuple
+
+class OrchidAccountError(Exception):
+    """Base class for Orchid account errors"""
+    pass
+
+class InvalidAddressError(OrchidAccountError):
+    """Invalid Ethereum address"""
+    pass
+
+class InvalidAmountError(OrchidAccountError):
+    """Invalid payment amount"""
+    pass
+
+class SigningError(OrchidAccountError):
+    """Error signing transaction or message"""
+    pass
+
+class OrchidAccount:
+    def __init__(self, 
+                 lottery: Lottery,
+                 funder_address: str,
+                 private_key: str):
+        try:
+            self.lottery = lottery
+            self.web3 = lottery.web3
+            self.funder = self.web3.to_checksum_address(funder_address)
+            self.key = private_key
+            self.signer = self.web3.eth.account.from_key(private_key).address
+        except ValueError as e:
+            raise InvalidAddressError(f"Invalid address format: {e}")
+        except Exception as e:
+            raise OrchidAccountError(f"Failed to initialize account: {e}")
+
+    def create_ticket(self, 
+                     amount: int,
+                     recipient: str,
+                     commitment: str,
+                     token_addr: str = "0x0000000000000000000000000000000000000000"
+                     ) -> str:
+        """
+        Create signed nanopayment ticket
+        
+        Args:
+            amount: Payment amount in wei
+            recipient: Recipient address
+            commitment: Random commitment hash
+            token_addr: Token contract address
+            
+        Returns:
+            Serialized ticket string
+        """
+        try:
+            if amount <= 0:
+                raise InvalidAmountError("Amount must be positive")
+
+            recipient = self.web3.to_checksum_address(recipient)
+            token_addr = self.web3.to_checksum_address(token_addr)
+
+            # Random nonce
+            nonce = secrets.randbits(128)
+
+            # Pack ticket data
+            packed0 = amount | (nonce << 128)
+            ratio = 0xffffffffffffffff  # Always create winning tickets for testing
+            packed1 = (ratio << 161) | (0 << 160)  # v=0
+
+            # Sign ticket
+            message_hash = self._get_ticket_hash(
+                token_addr,
+                recipient, 
+                commitment,
+                packed0,
+                packed1
+            )
+
+            sig = self.web3.eth.account.sign_message(
+                encode_defunct(message_hash),
+                private_key=self.key
+            )
+
+            # Adjust v and update packed1
+            v = sig.v - 27
+            packed1 = packed1 | v
+
+            # Format as hex strings
+            return (
+                hex(packed0)[2:].zfill(64) +
+                hex(packed1)[2:].zfill(64) +
+                hex(sig.r)[2:].zfill(64) +
+                hex(sig.s)[2:].zfill(64)
+            )
+
+        except OrchidAccountError:
+            raise
+        except Exception as e:
+            raise SigningError(f"Failed to create ticket: {e}")
+
+    def _get_ticket_hash(self,
+                        token_addr: str,
+                        recipient: str,
+                        commitment: str,
+                        packed0: int,
+                        packed1: int) -> bytes:
+        try:
+            return Web3.solidity_keccak(
+                ['bytes1', 'bytes1', 'address', 'bytes32', 'address', 'address', 
+                 'bytes32', 'uint256', 'uint256', 'bytes32'],
+                [b'\x19', b'\x00',
+                 self.lottery.contract_addr,
+                 b'\x00' * 31 + b'\x64',  # Chain ID
+                 token_addr,
+                 recipient,
+                 Web3.solidity_keccak(['bytes32'], [commitment]),
+                 packed0,
+                 packed1 >> 1,  # Remove v
+                 b'\x00' * 32]  # Empty data field
+            )
+        except Exception as e:
+            raise SigningError(f"Failed to create message hash: {e}")
+
+    async def get_balance(self,
+                         token_addr: str = "0x0000000000000000000000000000000000000000"
+                         ) -> Tuple[float, float]:
+        try:
+            balance, escrow = await self.lottery.check_balance(
+                token_addr,
+                self.funder,
+                self.signer
+            )
+            return (
+                self.lottery.wei_to_token(balance),
+                self.lottery.wei_to_token(escrow)
+            )
+        except Exception as e:
+            raise OrchidAccountError(f"Failed to get balance: {e}")

gai-backend/billing.py

@@ -1,33 +1,116 @@
 import json
+from redis.asyncio import Redis
+import redis
+from decimal import Decimal
+from typing import Optional, Dict
+import asyncio
 
-disconnect_threshold = -0.002
+class BillingError(Exception):
+    """Base class for billing errors that should terminate the connection"""
+    pass
 
-def invoice(amt):
-   return json.dumps({'type': 'invoice', 'amount': amt})
+class RedisConnectionError(BillingError):
+    """Redis connection or operation failed"""
+    pass
 
-class Billing:
-    def __init__(self, prices):
-       self.ledger = {}
-       self.prices = prices
-
-    def credit(self, id, type=None, amount=0):
-       self.adjust(id, type, amount, 1)
+class InconsistentStateError(BillingError):
+    """Billing state became inconsistent"""
+    pass
 
-    def debit(self, id, type=None, amount=0):
-       self.adjust(id, type, amount, -1)
-
-    def adjust(self, id, type, amount, sign):
-       amount_ = self.prices[type] if type is not None else amount
-       if id in self.ledger:
-          self.ledger[id] = self.ledger[id] + sign * amount_
-       else:
-          self.ledger[id] = sign * amount_
-
-    def min_balance(self):
-        return 2 * (self.prices['invoice'] + self.prices['payment'])
-
-    def balance(self, id):
-        if id in self.ledger:
-           return self.ledger[id]
-        else:
-            return 0
+class StrictRedisBilling:
+    def __init__(self, redis: Redis):
+        self.redis = redis
+
+    async def init(self):
+        try:
+            await self.redis.ping()
+        except Exception as e:
+            raise RedisConnectionError(f"Failed to connect to Redis: {e}")
+
+    def _get_client_key(self, client_id: str) -> str:
+        return f"billing:balance:{client_id}"
+
+    def _get_update_channel(self, client_id: str) -> str:
+        return f"billing:balance:updates:{client_id}"
+
+    async def credit(self, id: str, type: Optional[str] = None, amount: float = 0):
+        await self.adjust(id, type, amount, 1)
+
+    async def debit(self, id: str, type: Optional[str] = None, amount: float = 0):
+        await self.adjust(id, type, amount, -1)
+
+    async def adjust(self, id: str, type: Optional[str], amount: float, sign: int):
+        key = self._get_client_key(id)
+        channel = self._get_update_channel(id)
+
+        # Get amount from pricing if type is provided
+        amount_ = amount
+        if type is not None:
+            # Get price from config
+            config_data = await self.redis.get("config:data")
+            if not config_data:
+                raise BillingError("No configuration found")
+            config = json.loads(config_data)
+            price = config['billing']['prices'].get(type)
+            if price is None:
+                raise BillingError(f"Unknown price type: {type}")
+            amount_ = price
+
+        try:
+            async with self.redis.pipeline() as pipe:
+                while True:
+                    try:
+                        await pipe.watch(key)
+                        current = await self.redis.get(key)
+                        try:
+                            current_balance = Decimal(current) if current else Decimal('0')
+                        except (TypeError, ValueError) as e:
+                            raise InconsistentStateError(f"Invalid balance format in Redis: {e}")
+
+                        new_balance = current_balance + Decimal(str(sign * amount_))
+
+                        pipe.multi()
+                        await pipe.set(key, str(new_balance))
+                        await pipe.publish(channel, str(new_balance))
+                        await pipe.execute()
+                        return
+
+                    except redis.WatchError:
+                        continue
+
+                    except Exception as e:
+                        raise RedisConnectionError(f"Redis transaction failed: {e}")
+
+        except BillingError:
+            raise
+        except Exception as e:
+            raise RedisConnectionError(f"Unexpected Redis error: {e}")
+
+    async def balance(self, id: str) -> float:
+        try:
+            key = self._get_client_key(id)
+            balance = await self.redis.get(key)
+
+            if balance is None:
+                return 0
+
+            try:
+                return float(Decimal(balance))
+            except (TypeError, ValueError) as e:
+                raise InconsistentStateError(f"Invalid balance format: {e}")
+
+        except BillingError:
+            raise
+        except Exception as e:
+            raise RedisConnectionError(f"Failed to get balance: {e}")
+
+    async def min_balance(self) -> float:
+        try:
+            config_data = await self.redis.get("config:data")
+            if not config_data:
+                raise BillingError("No configuration found")
+            config = json.loads(config_data)
+            prices = config['billing']['prices']
+            return 2 * (prices['invoice'] + prices['payment'])
+        except Exception as e:
+            raise BillingError(f"Failed to calculate minimum balance: {e}")