tristero 0.1.7__py3-none-any.whl → 0.3.0__py3-none-any.whl

tristero/client.py

@@ -1,29 +1,61 @@
 import asyncio
+from dataclasses import dataclass
+from enum import Enum
 import json
 import logging
-from typing import Any, Optional, TypeVar
+import os
+import ssl
+from typing import Any, Dict, List, Literal, Optional, TypeVar, Union
 
 from eth_account.signers.local import LocalAccount
 from pydantic import BaseModel
 
-from tristero.api import ChainID, fill_order, poll_updates
-from .permit2 import create_order
+from tristero.api import (
+    fill_order,
+    fill_order_feather,
+    get_quote,
+    poll_updates,
+    poll_updates_feather,
+    get_margin_quote,
+    submit_margin_order,
+    submit_close_position,
+    get_margin_positions,
+    get_margin_position,
+    poll_margin_updates,
+)
+from .permit2 import (
+    Permit2Order,
+    create_permit2_order,
+    sign_margin_order,
+    sign_close_position,
+)
+from .data import ChainID
 from web3 import AsyncBaseProvider, AsyncWeb3
-import logging
-from web3 import AsyncWeb3
 from tenacity import (
     retry,
     stop_after_attempt,
     wait_exponential,
     retry_if_exception_type,
 )
+
+import certifi
+
 logger = logging.getLogger(__name__)
 
 P = TypeVar("P", bound=AsyncBaseProvider)
 
+
+class OrderType(str, Enum):
+    """Order type enum for unified API."""
+    SWAP = "swap"
+    FEATHER = "feather"
+    MARGIN = "margin"
+
+
 class WebSocketClosedError(Exception):
     pass
 
+
 class SwapException(Exception):
     """Base exception for all swap-related errors."""
     pass
@@ -41,80 +73,283 @@ class OrderFailedException(SwapException):
         self.order_id = order_id
         self.details = details
 
+
+class MarginException(Exception):
+    """Exception for margin-related errors."""
+    pass
+
+
 class TokenSpec(BaseModel, frozen=True):
     chain_id: ChainID
     token_address: str
 
-async def wait_for_completion(order_id: str):
+
+def make_async_w3(rpc_url: str) -> AsyncWeb3:
+    insecure_ssl = os.getenv("TRISTERO_INSECURE_SSL", "").strip().lower() in {"1", "true", "yes", "y", "on"}
+    if insecure_ssl:
+        provider = AsyncWeb3.AsyncHTTPProvider(rpc_url, request_kwargs={"ssl": False})
+        return AsyncWeb3(provider)
+
+    ca_file = os.getenv("TRISTERO_SSL_CA_FILE", "").strip()
+    if not ca_file:
+        ca_file = certifi.where()
+
+    if ca_file:
+        ssl_context = ssl.create_default_context(cafile=ca_file)
+        provider = AsyncWeb3.AsyncHTTPProvider(rpc_url, request_kwargs={"ssl": ssl_context})
+        return AsyncWeb3(provider)
+
+    provider = AsyncWeb3.AsyncHTTPProvider(rpc_url)
+    return AsyncWeb3(provider)
+
+
+@dataclass
+class QuoteResult:
+    """Result from get_quote operations."""
+    quote_data: Dict[str, Any]
+    order_type: str
+    
+    @property
+    def is_margin(self) -> bool:
+        return self.order_type == "MARGIN"
+    
+    @property
+    def is_feather(self) -> bool:
+        return self.order_type == "FEATHER"
+
+
+@dataclass 
+class OrderResult:
+    """Result from order submission."""
+    order_id: str
+    order_type: str
+    response: Dict[str, Any]
+
+
+@dataclass
+class MarginPosition:
+    """Margin position data."""
+    id: str
+    status: str
+    escrow_address: str
+    filler_address: str
+    taker_token_id: int
+    data: Dict[str, Any]
+
+
+async def wait_for_feather_completion(order_id: str):
+    ws = await poll_updates_feather(order_id)
+    try:
+        async for msg in ws:
+            if not msg:
+                continue
+            msg = json.loads(msg)
+            status = msg['status']
+            logger.info(
+                {
+                    "message": f"status={status}",
+                    "id": "order_update",
+                    "payload": msg,
+                }
+            )
+            if status in ['Expired']:
+                await ws.close()
+                raise SwapException(f"Swap failed: {ws.close_reason} {msg}")
+            elif status in ['Finalized']:
+                await ws.close()
+                return msg
+        raise WebSocketClosedError("WebSocket closed without completion status")
+    except Exception as e:
+        if not ws.close_code:
+            await ws.close()
+        raise
+
+
+async def wait_for_permit2_completion(order_id: str):
     ws = await poll_updates(order_id)
     try:
         async for msg in ws:
             msg = json.loads(msg)
+            if isinstance(msg, dict) and ("failed" in msg or "completed" in msg):
+                logger.info(
+                    {
+                        "message": f"failed={msg.get('failed')} completed={msg.get('completed')}",
+                        "id": "order_update",
+                        "payload": msg,
+                    }
+                )
+                if msg.get("failed"):
+                    await ws.close()
+                    raise SwapException(f"Swap failed: {ws.close_reason} {msg}")
+                elif msg.get("completed"):
+                    await ws.close()
+                    return msg
+                continue
+
+            if isinstance(msg, dict):
+                fill_tx = msg.get("fill_tx") or msg.get("fillTx")
+                if fill_tx:
+                    await ws.close()
+                    return msg
+
+                if msg.get("amount_repaid") is not None or msg.get("amount_settled") is not None:
+                    await ws.close()
+                    return msg
+
+            status = msg.get("status") if isinstance(msg, dict) else None
             logger.info(
                 {
-                    "message": f"failed={msg['failed']} completed={msg['completed']}",
+                    "message": f"status={status}",
                     "id": "order_update",
                     "payload": msg,
                 }
             )
-            if msg["failed"]:
+            if status in ["Failed", "Expired", "Rejected"]:
                 await ws.close()
-                raise Exception(f"Swap failed: {ws.close_reason} {msg}")
-            elif msg["completed"]:
+                raise SwapException(f"Swap failed: {ws.close_reason} {msg}")
+            elif status in ["Finalized", "Filled", "Completed"]:
                 await ws.close()
                 return msg
 
-        # If we exit the loop without completed/failed, raise to retry
         raise WebSocketClosedError("WebSocket closed without completion status")
     except Exception:
-        # Close cleanly if still open
         if not ws.close_code:
             await ws.close()
         raise
 
+
+async def wait_for_margin_completion(order_id: str):
+    """Wait for margin order completion via WebSocket."""
+    ws = await poll_margin_updates(order_id)
+    try:
+        async for msg in ws:
+            if not msg:
+                continue
+            msg = json.loads(msg)
+            if isinstance(msg, dict) and ("failed" in msg or "completed" in msg):
+                logger.info(
+                    {
+                        "message": f"failed={msg.get('failed')} completed={msg.get('completed')}",
+                        "id": "margin_order_update",
+                        "payload": msg,
+                    }
+                )
+                if msg.get("failed"):
+                    await ws.close()
+                    raise MarginException(f"Margin order failed: {msg}")
+                if msg.get("completed"):
+                    await ws.close()
+                    return msg
+                continue
+
+            status = msg.get('status', '') if isinstance(msg, dict) else ''
+            logger.info(
+                {
+                    "message": f"margin status={status}",
+                    "id": "margin_order_update",
+                    "payload": msg,
+                }
+            )
+            if status in ['Failed', 'Expired', 'Rejected']:
+                await ws.close()
+                raise MarginException(f"Margin order failed: {msg}")
+            elif status in ['Filled', 'Finalized', 'Completed']:
+                await ws.close()
+                return msg
+        raise WebSocketClosedError("WebSocket closed without completion status")
+    except Exception as e:
+        if not ws.close_code:
+            await ws.close()
+        raise
+
+
+async def wait_for_completion(order_id: str, order_type: Union[str, OrderType] = OrderType.SWAP):
+    """
+    Wait for order completion.
+    
+    Args:
+        order_id: Order ID to wait for
+        order_type: Type of order (swap, feather, margin)
+    """
+    if isinstance(order_type, str):
+        order_type = order_type.lower()
+    else:
+        order_type = order_type.value
+        
+    if order_type == "feather":
+        return await wait_for_feather_completion(order_id)
+    elif order_type == "margin":
+        return await wait_for_margin_completion(order_id)
+    else:
+        return await wait_for_permit2_completion(order_id)
+
+
 @retry(
     stop=stop_after_attempt(3),
     wait=wait_exponential(multiplier=1, min=4, max=10),
     retry=retry_if_exception_type(WebSocketClosedError)
 )
-async def wait_for_completion_with_retry(order_id: str):
-    return await wait_for_completion(order_id)
+async def wait_for_completion_with_retry(
+    order_id: str, 
+    order_type: Union[str, OrderType] = OrderType.SWAP
+):
+    """Wait for order completion with automatic retry on WebSocket failures."""
+    return await wait_for_completion(order_id, order_type)
 
-async def start_swap(w3: AsyncWeb3[P], account: LocalAccount, from_t: TokenSpec, to_t: TokenSpec, raw_amount: int) -> str:
-    """
-    Execute a token swap operation.
 
-    Args:
-        w3: Web3 provider instance
-        account: Account to execute swap from
-        from_t: Source token specification
-        to_t: Target token specification
-        raw_amount: Amount in smallest unit (e.g., wei)
+@dataclass
+class FeatherSwapResult:
+    deposit_address: str
+    data: Any
 
-    Returns:
-        Order ID for tracking the swap
 
-    Raises:
-        Exception: If order creation or submission fails
-    """
-    data, sig = await create_order(
+class FeatherException(Exception):
+    pass
+
+
+async def start_feather_swap(
+    src_t: TokenSpec,
+    dst_t: TokenSpec,
+    dst_addr: str,
+    raw_amount: int,
+    client_id: str = ''
+):
+    resp = await fill_order_feather(str(src_t.chain_id.value), str(dst_t.chain_id.value), dst_addr, raw_amount, client_id)
+    if resp['detail']:
+        raise FeatherException(resp)
+    else:
+        return FeatherSwapResult(
+            deposit_address=resp['deposit_address'],
+            data=resp
+        )
+
+
+async def start_permit2_swap(
+    w3: AsyncWeb3[P],
+    account: LocalAccount,
+    src_t: TokenSpec,
+    dst_t: TokenSpec,
+    raw_amount: int,
+):
+    order = await create_permit2_order(
         w3,
         account,
-        from_t.chain_id,
-        from_t.token_address,
-        to_t.chain_id,
-        to_t.token_address,
+        str(src_t.chain_id.value),
+        src_t.token_address,
+        str(dst_t.chain_id.value),
+        dst_t.token_address,
         raw_amount,
     )
     response = await fill_order(
-        str(sig.signature.to_0x_hex()),
-        data.domain.model_dump(by_alias=True, mode="json"),
-        data.message.model_dump(by_alias=True, mode="json"),
+        str(order.sig.signature.to_0x_hex()),
+        order.msg.domain.model_dump(by_alias=True, mode="json"),
+        order.msg.message.model_dump(by_alias=True, mode="json"),
     )
+    order_id = response['id']
+    return order_id
 
-    return response['id']
 
-async def execute_swap(
+async def execute_permit2_swap(
     w3: AsyncWeb3[P],
     account: LocalAccount,
     src_t: TokenSpec,
@@ -124,24 +359,310 @@ async def execute_swap(
     timeout: Optional[float] = None
 ) -> dict[str, Any]:
     """Execute and wait for swap completion."""
-    order_id = await start_swap(
-        w3,
-        account,
-        src_t,
-        dst_t,
-        raw_amount
-    )
+    order_id = await start_permit2_swap(w3, account, src_t, dst_t, raw_amount)
     logger.info(f"Swap order placed: {order_id}")
 
     waiter = wait_for_completion_with_retry if retry else wait_for_completion
 
     try:
         if timeout is None:
-            return await waiter(order_id)
+            return await waiter(order_id, OrderType.SWAP)
 
         return await asyncio.wait_for(
-            waiter(order_id),
+            waiter(order_id, OrderType.SWAP),
             timeout=timeout
         )
     except asyncio.TimeoutError as exc:
         raise StuckException(f"Swap {order_id} timed out after {timeout}s") from exc
+
+
+async def execute_swap(
+    w3: AsyncWeb3[P],
+    account: LocalAccount,
+    src_t: TokenSpec,
+    dst_t: TokenSpec,
+    raw_amount: int,
+    retry: bool = True,
+    timeout: Optional[float] = None,
+) -> dict[str, Any]:
+    return await execute_permit2_swap(
+        w3=w3,
+        account=account,
+        src_t=src_t,
+        dst_t=dst_t,
+        raw_amount=raw_amount,
+        retry=retry,
+        timeout=timeout,
+    )
+
+
+async def request_margin_quote(
+    chain_id: str,
+    wallet_address: str,
+    quote_currency: str,
+    base_currency: str,
+    leverage_ratio: int,
+    collateral_amount: str,
+) -> QuoteResult:
+    """
+    Request a quote for opening a margin position.
+    
+    Args:
+        chain_id: Chain ID (e.g., "42161" for Arbitrum)
+        wallet_address: User's wallet address
+        quote_currency: Quote currency token address (e.g., USDC)
+        base_currency: Base currency token address (e.g., WETH)
+        leverage_ratio: Leverage ratio (e.g., 2 for 2x leverage)
+        collateral_amount: Collateral amount in raw units
+    
+    Returns:
+        QuoteResult containing quote data
+    """
+    quote = await get_margin_quote(
+        chain_id=chain_id,
+        wallet_address=wallet_address,
+        quote_currency=quote_currency,
+        base_currency=base_currency,
+        leverage_ratio=leverage_ratio,
+        collateral_amount=collateral_amount,
+    )
+    return QuoteResult(
+        quote_data=quote,
+        order_type="MARGIN",
+    )
+
+
+def sign_order(
+    quote: QuoteResult,
+    private_key: str,
+) -> Dict[str, Any]:
+    """
+    Sign a quote to create a submittable order.
+    
+    Args:
+        quote: Quote result from request_margin_quote
+        private_key: Private key for signing
+    
+    Returns:
+        Signed order payload
+    """
+    if quote.is_margin:
+        return sign_margin_order(quote.quote_data, private_key)
+    else:
+        raise ValueError(f"Unsupported order type for signing: {quote.order_type}")
+
+
+async def submit_order(signed_order: Dict[str, Any]) -> OrderResult:
+    """
+    Submit a signed order.
+    
+    Args:
+        signed_order: Signed order from sign_order
+    
+    Returns:
+        OrderResult with order_id
+    """
+    response = await submit_margin_order(signed_order)
+    return OrderResult(
+        order_id=response.get("id", response.get("order_id", "")),
+        order_type="MARGIN",
+        response=response,
+    )
+
+
+async def open_margin_position(
+    private_key: str,
+    chain_id: str,
+    wallet_address: str,
+    quote_currency: str,
+    base_currency: str,
+    leverage_ratio: int,
+    collateral_amount: str,
+    wait_for_result: bool = True,
+    timeout: Optional[float] = None,
+) -> Dict[str, Any]:
+    """
+    Open a margin position in one call (quote + sign + submit + optionally wait).
+    
+    Args:
+        private_key: Private key for signing
+        chain_id: Chain ID
+        wallet_address: User's wallet address
+        quote_currency: Quote currency token address
+        base_currency: Base currency token address
+        leverage_ratio: Leverage ratio
+        collateral_amount: Collateral amount in raw units
+        wait_for_result: Whether to wait for order completion
+        timeout: Optional timeout in seconds
+    
+    Returns:
+        Order result or completion result
+    """
+    quote = await request_margin_quote(
+        chain_id=chain_id,
+        wallet_address=wallet_address,
+        quote_currency=quote_currency,
+        base_currency=base_currency,
+        leverage_ratio=leverage_ratio,
+        collateral_amount=collateral_amount,
+    )
+    
+    signed = sign_order(quote, private_key)
+    result = await submit_order(signed)
+    
+    logger.info(f"Margin order submitted: {result.order_id}")
+    
+    if not wait_for_result:
+        return {
+            "order_id": result.order_id,
+            "quote": quote.quote_data,
+            "response": result.response,
+        }
+    
+    try:
+        if timeout:
+            completion = await asyncio.wait_for(
+                wait_for_completion_with_retry(result.order_id, OrderType.MARGIN),
+                timeout=timeout
+            )
+        else:
+            completion = await wait_for_completion_with_retry(result.order_id, OrderType.MARGIN)
+        
+        return {
+            "order_id": result.order_id,
+            "quote": quote.quote_data,
+            "response": result.response,
+            "completion": completion,
+        }
+    except asyncio.TimeoutError as exc:
+        raise StuckException(f"Margin order {result.order_id} timed out after {timeout}s") from exc
+
+
+async def close_margin_position(
+    private_key: str,
+    chain_id: str,
+    position_id: int,
+    escrow_contract: str,
+    authorized: str,
+    cash_settle: bool = False,
+    fraction_bps: int = 10_000,
+    deadline_seconds: int = 3600,
+    wait_for_result: bool = True,
+    timeout: Optional[float] = None,
+) -> Dict[str, Any]:
+    """
+    Close a margin position.
+    
+    Args:
+        private_key: Private key for signing
+        chain_id: Chain ID
+        position_id: Position ID (NFT token ID / taker_token_id)
+        escrow_contract: Escrow contract address
+        authorized: Authorized filler address
+        cash_settle: Whether to cash settle (True) or swap settle (False)
+        fraction_bps: Fraction to close in basis points (10000 = 100%)
+        deadline_seconds: Deadline in seconds from now
+        wait_for_result: Whether to wait for completion
+        timeout: Optional timeout in seconds
+    
+    Returns:
+        Close result
+    """
+    signed = sign_close_position(
+        chain_id=int(chain_id),
+        position_id=position_id,
+        private_key=private_key,
+        escrow_contract=escrow_contract,
+        authorized=authorized,
+        cash_settle=cash_settle,
+        fraction_bps=fraction_bps,
+        deadline_seconds=deadline_seconds,
+    )
+    
+    response = await submit_close_position(signed)
+    close_order = response.get("close_order")
+    close_order_id = None
+    if isinstance(close_order, dict):
+        close_order_id = close_order.get("id")
+
+    order_id = (
+        close_order_id
+        or response.get("id")
+        or response.get("order_id")
+        or response.get("orderId")
+        or response.get("orderID")
+        or ""
+    )
+    if not isinstance(order_id, str) or not order_id:
+        raise MarginException(f"Close position response missing order id: {response}")
+    
+    logger.info(f"Close position order submitted: {order_id}")
+    
+    if not wait_for_result:
+        return {
+            "order_id": order_id,
+            "response": response,
+        }
+    
+    try:
+        if timeout:
+            completion = await asyncio.wait_for(
+                wait_for_completion_with_retry(order_id, OrderType.SWAP),
+                timeout=timeout
+            )
+        else:
+            completion = await wait_for_completion_with_retry(order_id, OrderType.SWAP)
+        
+        return {
+            "order_id": order_id,
+            "response": response,
+            "completion": completion,
+        }
+    except asyncio.TimeoutError as exc:
+        raise StuckException(f"Close position {order_id} timed out after {timeout}s") from exc
+
+
+async def list_margin_positions(wallet_address: str) -> List[MarginPosition]:
+    """
+    Get all margin positions for a wallet.
+    
+    Args:
+        wallet_address: User's wallet address
+    
+    Returns:
+        List of MarginPosition objects
+    """
+    positions = await get_margin_positions(wallet_address)
+    return [
+        MarginPosition(
+            id=str(p.get("id", "")),
+            status=p.get("status", ""),
+            escrow_address=p.get("escrow_address", ""),
+            filler_address=p.get("filler_address", ""),
+            taker_token_id=p.get("taker_token_id", 0),
+            data=p,
+        )
+        for p in positions
+    ]
+
+
+async def get_position(position_id: str) -> MarginPosition:
+    """
+    Get a specific margin position.
+    
+    Args:
+        position_id: Position ID
+    
+    Returns:
+        MarginPosition object
+    """
+    p = await get_margin_position(position_id)
+    return MarginPosition(
+        id=str(p.get("id", "")),
+        status=p.get("status", ""),
+        escrow_address=p.get("escrow_address", ""),
+        filler_address=p.get("filler_address", ""),
+        taker_token_id=p.get("taker_token_id", 0),
+        data=p,
+    )
+