uni-exec-engine 0.2.0__py3-none-any.whl

uniswap_autopilot/state_machine.py

@@ -0,0 +1,315 @@
+"""Trade execution state machine — idempotent, resumable, auditable.
+
+Provides a lightweight checkpoint mechanism for transaction workflows that
+share the same ``run_id`` used by the audit system.  Each state transition is
+recorded to a JSON state file so that a crashed orchestrator can resume from
+the last safe checkpoint without re-signing or re-broadcasting transactions.
+
+States are the same event names used by ``audit.py`` so the audit log and the
+state machine speak the same language:
+
+    INIT  →  PREFLIGHT  →  SIGNED  →  BROADCAST  →  CONFIRMED
+
+Any state can transition to FAILED (terminal).  CANCELLED is also terminal.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import tempfile
+import threading
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+# ── Constants ───────────────────────────────────────────────────────────────
+
+STATE_INIT = "init"
+STATE_PREFLIGHT = "preflight"
+STATE_SIGNED = "signed"
+STATE_BROADCAST = "broadcast"
+STATE_CONFIRMED = "confirmed"
+STATE_CANCELLED = "cancelled"
+STATE_FAILED = "failed"
+
+TERMINAL_STATES: frozenset[str] = frozenset({STATE_CONFIRMED, STATE_CANCELLED, STATE_FAILED})
+
+# Ordered so the comparison STATE_ORDER.index(a) < STATE_ORDER.index(b) works.
+STATE_ORDER: tuple[str, ...] = (
+    STATE_INIT,
+    STATE_PREFLIGHT,
+    STATE_SIGNED,
+    STATE_BROADCAST,
+    STATE_CONFIRMED,
+)
+
+# Valid transition map — any state can additionally go to FAILED.
+_VALID_TRANSITIONS: dict[str, frozenset[str]] = {
+    STATE_INIT: frozenset({STATE_PREFLIGHT, STATE_CANCELLED}),
+    STATE_PREFLIGHT: frozenset({STATE_SIGNED, STATE_CANCELLED}),
+    STATE_SIGNED: frozenset({STATE_BROADCAST, STATE_CANCELLED}),
+    STATE_BROADCAST: frozenset({STATE_CONFIRMED}),
+}
+
+
+# ── Module-level state ──────────────────────────────────────────────────────
+
+_write_lock = threading.Lock()
+_DEFAULT_PROJECT = __name__.split(".")[0].replace("_", "-")
+
+
+def _state_dir() -> Path:
+    """Return the directory used for state machine checkpoint files.
+
+    The directory can be overridden with the ``STAGEFORGE_STATE_DIR``
+    environment variable; otherwise ``~/.stageforge/states`` is used.
+    """
+    env = os.environ.get("STAGEFORGE_STATE_DIR", "").strip()
+    if env:
+        return Path(env)
+    return Path.home() / ".stageforge" / "states"
+
+
+def _state_path(run_id: str) -> Path:
+    """File-system path for a given *run_id*."""
+    return _state_dir() / f"{run_id}.json"
+
+
+def _now_iso() -> str:
+    """ISO-8601 UTC timestamp (without microsecond noise)."""
+    return datetime.now(tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
+
+
+def _resolve_run_id(explicit: str | None) -> str | None:
+    if explicit:
+        return explicit
+    for name in ("STAGEFORGE_RUN_ID", "AUDIT_RUN_ID", "RUN_ID"):
+        value = (os.environ.get(name) or "").strip()
+        if value:
+            return value
+    return None
+
+
+# ── Action mapping (happy-path states → next action) ─────────────────────
+
+_STATE_TO_ACTION: dict[str, str | None] = {
+    STATE_INIT: STATE_PREFLIGHT,
+    STATE_PREFLIGHT: STATE_SIGNED,
+    STATE_SIGNED: STATE_BROADCAST,
+    STATE_BROADCAST: STATE_CONFIRMED,
+    STATE_CONFIRMED: None,
+    STATE_CANCELLED: None,
+    STATE_FAILED: None,
+}
+
+
+# ── Public API ───────────────────────────────────────────────────────────────
+
+
+def load_state(run_id: str) -> dict[str, Any] | None:
+    """Read an existing state checkpoint, or ``None`` if not found."""
+    path = _state_path(run_id)
+    try:
+        with path.open() as fh:
+            return json.load(fh)  # type: ignore[no-any-return]
+    except (FileNotFoundError, json.JSONDecodeError):
+        return None
+
+
+def init_state(
+    run_id: str,
+    *,
+    project: str | None = None,
+    payload: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    """Create (or return) the initial state record.
+
+    If a state file already exists for *run_id* it is returned unchanged
+    (idempotent).  Callers that want to start over must delete the file
+    first.
+    """
+    existing = load_state(run_id)
+    if existing is not None:
+        return existing
+
+    now = _now_iso()
+    state: dict[str, Any] = {
+        "run_id": run_id,
+        "project": project or _DEFAULT_PROJECT,
+        "current_state": STATE_INIT,
+        "created_at": now,
+        "updated_at": now,
+        "transition_log": [],
+        "payload": payload or {},
+    }
+    _atomic_write(_state_path(run_id), state)
+    return state
+
+
+def transition(
+    run_id: str,
+    target_state: str,
+    *,
+    payload: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    """Advance the state machine to *target_state*.
+
+    Rules:
+
+    * If *target_state* is FAILED the transition is always accepted
+      (any state can fail).
+    * If the current state is already >= *target_state* in the happy-path
+      order, the call is a no-op (idempotent).
+    * If the current state is a terminal state, raises ``RuntimeError``.
+    * If the transition is not allowed, raises ``ValueError``.
+
+    Returns the *updated* state dict.
+    """
+    state = init_state(run_id)  # ensure state file exists
+    current = state["current_state"]
+
+    # Terminal state — no-op if already in target, error otherwise.
+    if current in TERMINAL_STATES:
+        if current == target_state:
+            return state
+        raise RuntimeError(
+            f"State machine for {run_id} is in terminal state {current!r} — cannot transition to {target_state!r}."
+        )
+
+    # Idempotent — already at or beyond this state.
+    if (
+        target_state != STATE_FAILED
+        and current in STATE_ORDER
+        and target_state in STATE_ORDER
+        and STATE_ORDER.index(current) >= STATE_ORDER.index(target_state)
+    ):
+        return state
+
+    # Validate transition.
+    if target_state == STATE_FAILED:
+        pass  # always allowed
+    else:
+        allowed = _VALID_TRANSITIONS.get(current, frozenset())
+        if target_state not in allowed:
+            raise ValueError(
+                f"Illegal transition for {run_id}: {current!r} → {target_state!r}. Allowed: {sorted(allowed)}"
+            )
+
+    # Perform transition.
+    now = _now_iso()
+    transition_entry: dict[str, Any] = {
+        "from": current,
+        "to": target_state,
+        "at": now,
+    }
+    state["current_state"] = target_state
+    state["updated_at"] = now
+    state["transition_log"].append(transition_entry)
+    if payload is not None:
+        state["payload"].update(payload)
+
+    _atomic_write(_state_path(run_id), state)
+    return state
+
+
+def status(run_id: str) -> dict[str, Any]:
+    """Return a high-level status summary for *run_id*."""
+    state = load_state(run_id)
+    if state is None:
+        return {"run_id": run_id, "found": False, "current_state": None}
+    return {
+        "run_id": run_id,
+        "found": True,
+        "project": state.get("project"),
+        "current_state": state.get("current_state"),
+        "created_at": state.get("created_at"),
+        "updated_at": state.get("updated_at"),
+        "transition_count": len(state.get("transition_log", [])),
+    }
+
+
+def next_action(run_id: str) -> str | None:
+    """Return the next allowed happy-path action for *run_id*, or None.
+
+    Callers should check this before executing any side effect (sign, broadcast,
+    confirm).  If the returned action does not match what the caller plans to do
+    the caller must skip that step — it has already been completed in a prior
+    run (anti-replay).
+    """
+    state = load_state(run_id)
+    if state is None:
+        return STATE_PREFLIGHT  # fresh run — start at preflight
+    return _STATE_TO_ACTION.get(state.get("current_state", STATE_INIT))
+
+
+def resume(run_id: str) -> str | None:
+    """Return the action that an interrupted run should resume from, or None.
+
+    ``resume()`` is a convenience wrapper around ``next_action()`` with clearer
+    semantics for orchestrators.
+
+    Returns ``None`` when the run is already complete (CONFIRMED / CANCELLED /
+    FAILED) or the state file is missing.
+    """
+    return next_action(run_id)
+
+
+def list_runs() -> list[dict[str, Any]]:
+    """Return status summaries for all persisted runs."""
+    sd = _state_dir()
+    results: list[dict[str, Any]] = []
+    if not sd.is_dir():
+        return results
+    for p in sorted(sd.glob("*.json")):
+        run_id = p.stem
+        s = load_state(run_id)
+        if s is not None:
+            results.append(
+                {
+                    "run_id": run_id,
+                    "project": s.get("project"),
+                    "current_state": s.get("current_state"),
+                    "updated_at": s.get("updated_at"),
+                }
+            )
+    return results
+
+
+# ── Internal helpers ─────────────────────────────────────────────────────────
+
+
+def _atomic_write(path: Path, data: dict[str, Any]) -> None:
+    """Write *data* to *path* atomically (write-tmp + rename)."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with _write_lock:
+        fd, tmp = tempfile.mkstemp(dir=str(path.parent), suffix=".tmp")
+        try:
+            with os.fdopen(fd, "w", encoding="utf-8") as fh:
+                json.dump(data, fh, ensure_ascii=False, indent=2)
+            os.replace(tmp, path)
+        except Exception:
+            try:
+                os.unlink(tmp)
+            except OSError:
+                pass
+            raise
+
+
+__all__ = [
+    "STATE_BROADCAST",
+    "STATE_CANCELLED",
+    "STATE_CONFIRMED",
+    "STATE_FAILED",
+    "STATE_INIT",
+    "STATE_PREFLIGHT",
+    "STATE_SIGNED",
+    "TERMINAL_STATES",
+    "init_state",
+    "list_runs",
+    "load_state",
+    "next_action",
+    "resume",
+    "status",
+    "transition",
+]

uniswap_autopilot/swap/__init__.py

@@ -0,0 +1 @@
+"""Swap modules: trading API, flow execution, and extensions."""

uniswap_autopilot/swap/deep_link.py

@@ -0,0 +1,69 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+
+
+from uniswap_autopilot.common.common import (
+    add_common_arguments,
+    build_swap_link,
+    dump_json,
+    normalize_chain,
+    override_decimals,
+    parse_amount,
+    resolve_token,
+)
+
+
+def build_swap_link_response(
+    chain_name: str,
+    token_in_name: str,
+    token_out_name: str,
+    amount_value: str,
+    field: str = "INPUT",
+    token_in_decimals: int | None = None,
+    token_out_decimals: int | None = None,
+) -> dict[str, object]:
+    chain = normalize_chain(chain_name)
+    amount = parse_amount(amount_value)
+    token_in = override_decimals(resolve_token(chain, token_in_name), token_in_decimals)
+    token_out = override_decimals(resolve_token(chain, token_out_name), token_out_decimals)
+    human_amount = format(amount, "f")
+    link = build_swap_link(chain, token_in, token_out, human_amount, field)
+    return {
+        "action": "build_swap_link",
+        "chain": {"key": chain.key, "chainId": chain.chain_id},
+        "tokenIn": token_in,
+        "tokenOut": token_out,
+        "amount": human_amount,
+        "field": field,
+        "deepLink": link,
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="生成 Uniswap swap deep link")
+    add_common_arguments(parser)
+    parser.add_argument(
+        "--field",
+        choices=["INPUT", "OUTPUT"],
+        default="INPUT",
+        help="value 对应输入还是输出字段",
+    )
+    args = parser.parse_args()
+
+    dump_json(
+        build_swap_link_response(
+            chain_name=args.chain,
+            token_in_name=args.token_in,
+            token_out_name=args.token_out,
+            amount_value=args.amount,
+            field=args.field,
+            token_in_decimals=args.token_in_decimals,
+            token_out_decimals=args.token_out_decimals,
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()

uniswap_autopilot/swap/extensions/__init__.py

@@ -0,0 +1,2 @@
+"""Optional swap extension commands."""
+

uniswap_autopilot/swap/extensions/bridge.py

@@ -0,0 +1,197 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import time
+from typing import Any
+
+
+from uniswap_autopilot.common.common import (
+    dump_json,
+    load_local_env,
+    normalize_chain,
+    override_decimals,
+    parse_amount,
+    post_json,
+    require_api_key,
+    resolve_api_token,
+    resolve_token,
+    resolve_wallet_address,
+    decimal_to_base_units,
+    native_currency_address,
+)
+
+
+def build_cross_chain_quote_payload(
+    wallet: str | None,
+    src_chain_id: int,
+    dst_chain_id: int,
+    api_token_in: dict[str, Any],
+    api_token_out: dict[str, Any],
+    base_amount: str,
+    swap_type: str,
+    slippage: float,
+) -> dict[str, Any]:
+    return {
+        "swapper": wallet or native_currency_address(),
+        "tokenIn": api_token_in["address"],
+        "tokenOut": api_token_out["address"],
+        "tokenInChainId": src_chain_id,
+        "tokenOutChainId": dst_chain_id,
+        "amount": base_amount,
+        "type": swap_type,
+        "slippageTolerance": slippage,
+    }
+
+
+def prepare_cross_chain_quote(
+    src_chain_name: str,
+    dst_chain_name: str,
+    token_in_name: str,
+    token_out_name: str,
+    amount_value: str,
+    wallet: str | None = None,
+    token_in_decimals: int | None = None,
+    token_out_decimals: int | None = None,
+    swap_type: str = "EXACT_INPUT",
+    slippage: float = 0.5,
+) -> dict[str, Any]:
+    src_chain = normalize_chain(src_chain_name)
+    dst_chain = normalize_chain(dst_chain_name)
+    amount = parse_amount(amount_value)
+    token_in = override_decimals(resolve_token(src_chain, token_in_name), token_in_decimals)
+    token_out = override_decimals(resolve_token(dst_chain, token_out_name), token_out_decimals)
+    api_token_in = resolve_api_token(src_chain, token_in)
+    api_token_out = resolve_api_token(dst_chain, token_out)
+    validated_wallet = resolve_wallet_address(wallet, "wallet")
+    base_amount = decimal_to_base_units(amount, api_token_in["decimals"])
+
+    payload = build_cross_chain_quote_payload(
+        wallet=validated_wallet,
+        src_chain_id=src_chain.chain_id,
+        dst_chain_id=dst_chain.chain_id,
+        api_token_in=api_token_in,
+        api_token_out=api_token_out,
+        base_amount=base_amount,
+        swap_type=swap_type,
+        slippage=slippage,
+    )
+
+    response: dict[str, Any] = {
+        "action": "cross_chain_quote",
+        "srcChain": {"key": src_chain.key, "chainId": src_chain.chain_id},
+        "dstChain": {"key": dst_chain.key, "chainId": dst_chain.chain_id},
+        "tokenIn": token_in,
+        "tokenOut": token_out,
+        "humanAmount": format(amount, "f"),
+        "baseAmount": base_amount,
+        "requestPayload": payload,
+    }
+    return response
+
+
+def fetch_cross_chain_quote(response: dict[str, Any]) -> dict[str, Any]:
+    api_key = require_api_key()
+    raw = post_json("quote", response["requestPayload"], api_key)
+    response["rawQuote"] = raw
+    return response
+
+
+def verify_bridge_arrival(
+    dst_chain_name: str,
+    token_address: str,
+    wallet: str,
+    expected_amount_base: str,
+    timeout: int = 600,
+    poll_interval: int = 30,
+    rpc_url: str | None = None,
+) -> dict[str, Any]:
+    from uniswap_autopilot.common.common import PUBLIC_RPC_URLS
+    from uniswap_autopilot.execute._internal.rpc import query_erc20_balance, query_native_balance, resolve_rpc_url
+
+    dst_chain = normalize_chain(dst_chain_name)
+    rpc = rpc_url or PUBLIC_RPC_URLS.get(dst_chain.key) or resolve_rpc_url(None, dst_chain.chain_id)[0]
+    if not rpc:
+        raise RuntimeError(f"RPC URL not configured for {dst_chain_name}")
+
+    expected = int(expected_amount_base)
+    start = time.time()
+    polls = 0
+    balance = 0
+
+    while time.time() - start < timeout:
+        polls += 1
+        try:
+            if token_address in ("NATIVE", native_currency_address()):
+                balance = query_native_balance(wallet, rpc)
+            else:
+                balance = query_erc20_balance(wallet, token_address, rpc)
+            if balance >= expected:
+                return {
+                    "arrived": True,
+                    "balance": str(balance),
+                    "expected": str(expected),
+                    "elapsedSeconds": int(time.time() - start),
+                    "polls": polls,
+                }
+        except Exception:
+            pass
+        time.sleep(poll_interval)
+
+    return {
+        "arrived": False,
+        "balance": str(balance),
+        "expected": str(expected),
+        "elapsedSeconds": int(time.time() - start),
+        "polls": polls,
+        "timeout": True,
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Cross-chain swap via Uniswap Trading API")
+    sub = parser.add_subparsers(dest="command")
+
+    q = sub.add_parser("quote", help="Get cross-chain quote")
+    q.add_argument("--src-chain", required=True, help="Source chain")
+    q.add_argument("--dst-chain", required=True, help="Destination chain")
+    q.add_argument("--token-in", required=True)
+    q.add_argument("--token-out", required=True)
+    q.add_argument("--amount", required=True)
+    q.add_argument("--wallet")
+    q.add_argument("--slippage", type=float, default=0.5)
+    q.add_argument("--request-only", action="store_true", help="Only print payload")
+
+    v = sub.add_parser("verify", help="Verify bridge arrival")
+    v.add_argument("--dst-chain", required=True)
+    v.add_argument("--token-address", required=True)
+    v.add_argument("--wallet", required=True)
+    v.add_argument("--expected-amount", required=True, help="Expected amount in base units")
+    v.add_argument("--timeout", type=int, default=600)
+    v.add_argument("--poll-interval", type=int, default=30)
+
+    args = parser.parse_args()
+    load_local_env()
+
+    if args.command == "quote":
+        result = prepare_cross_chain_quote(
+            args.src_chain, args.dst_chain, args.token_in, args.token_out,
+            args.amount, args.wallet, slippage=args.slippage,
+        )
+        if not args.request_only:
+            result = fetch_cross_chain_quote(result)
+        dump_json(result)
+    elif args.command == "verify":
+        result = verify_bridge_arrival(
+            args.dst_chain, args.token_address, args.wallet,
+            args.expected_amount, args.timeout, args.poll_interval,
+        )
+        status = "ARRIVED" if result["arrived"] else "TIMEOUT"
+        print(f"Bridge status: {status} ({result['polls']} polls, {result['elapsedSeconds']}s)")
+        dump_json(result)
+    else:
+        parser.print_help()
+
+
+if __name__ == "__main__":
+    main()